Integrate with client, cli, app, mcp

haiku_rag_slim/haiku/rag/agents/rlm/__init__.py

@@ -1,12 +1,11 @@
 from haiku.rag.agents.rlm.agent import create_rlm_agent
-from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext, RLMDeps
+from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
 from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
 from haiku.rag.agents.rlm.prompts import RLM_SYSTEM_PROMPT
 from haiku.rag.agents.rlm.sandbox import REPLEnvironment, REPLResult
 
 __all__ = [
     "CodeExecution",
-    "RLMConfig",
     "RLMContext",
     "RLMDeps",
     "RLMResult",

haiku_rag_slim/haiku/rag/agents/rlm/agent.py

@@ -16,7 +16,7 @@ def _get_or_create_repl(ctx) -> REPLEnvironment:
     if key not in _repl_cache:
         _repl_cache[key] = REPLEnvironment(
             client=ctx.deps.client,
-            config=ctx.deps.rlm_config,
+            config=ctx.deps.config.rlm,
             context=ctx.deps.context,
         )
     return _repl_cache[key]
@@ -35,7 +35,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
     Returns:
         A pydantic-ai Agent configured for RLM execution.
     """
-    model = get_model(config.qa.model, config)
+    model = get_model(config.rlm.model, config)
 
     agent: Agent[RLMDeps, RLMResult] = Agent(  # type: ignore[invalid-assignment]
         model,

haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py

@@ -1,8 +1,6 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel
-
 from haiku.rag.store.models import Document, SearchResult
 
 if TYPE_CHECKING:
@@ -11,19 +9,12 @@ if TYPE_CHECKING:
     from haiku.rag.config.models import AppConfig
 
 
-class RLMConfig(BaseModel):
-    """Configuration for RLM agent sandbox execution."""
-
-    code_timeout: float = 60.0
-    max_output_chars: int = 50_000
-    max_tool_calls: int = 20
-
-
 @dataclass
 class RLMContext:
     """Mutable context accumulating data during RLM execution."""
 
     documents: list[Document] | None = None
+    filter: str | None = None
     search_results: list[SearchResult] = field(default_factory=list)
     code_executions: "list[CodeExecution]" = field(default_factory=list)
 
@@ -34,5 +25,4 @@ class RLMDeps:
 
     client: "HaikuRAG"
     config: "AppConfig"
-    rlm_config: RLMConfig = field(default_factory=RLMConfig)
     context: RLMContext = field(default_factory=RLMContext)

haiku_rag_slim/haiku/rag/agents/rlm/models.py

@@ -1,7 +1,5 @@
 from pydantic import BaseModel, Field
 
-from haiku.rag.agents.research.models import Citation
-
 
 class CodeExecution(BaseModel):
     """Result of executing a code block in the RLM sandbox."""
@@ -16,10 +14,6 @@ class RLMResult(BaseModel):
     """Result from RLM agent execution."""
 
     answer: str = Field(description="The answer to the user's question")
-    citations: list[Citation] = Field(
-        default_factory=list,
-        description="Citations for sources used in the answer",
-    )
     code_executions: list[CodeExecution] = Field(
         default_factory=list,
         description="History of code executions during the RLM session",

haiku_rag_slim/haiku/rag/agents/rlm/prompts.py

@@ -1,14 +1,20 @@
 RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-You have access to a sandboxed Python environment with these haiku.rag functions:
+IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
+
+CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+- search("query")  ✓ CORRECT
+- from haiku.rag import search  ✗ WRONG - will fail
+
+You have access to a sandboxed Python environment with these haiku.rag functions (use them directly, no imports needed):
 
 ## Available Functions
 
-### search(query, limit=10, filter=None) -> list[dict]
+### search(query, limit=10) -> list[dict]
 Search the knowledge base using hybrid search (vector + full-text).
 Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-### list_documents(limit=10, offset=0, filter=None) -> list[dict]
+### list_documents(limit=10, offset=0) -> list[dict]
 List available documents in the knowledge base.
 Returns list of dicts with keys: id, title, uri, created_at
 
@@ -21,7 +27,7 @@ Get the structured DoclingDocument object for advanced analysis.
 Returns a DoclingDocument object, or None if not found.
 See "DoclingDocument API" section below for how to use it.
 
-### ask(question, filter=None) -> str
+### ask(question) -> str
 Ask a question using the QA agent with RAG. Returns the answer as a string.
 Use this for semantic analysis that benefits from LLM reasoning.
 
@@ -30,12 +36,13 @@ You can import: json, re, collections, math, statistics, itertools, functools, d
 
 ## Strategy Guide
 
-1. **Explore First**: Start by listing documents or searching to understand what's available.
+1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-2. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
+2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
-3. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with.
+3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
-4. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections.
+4. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with.
-5. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function.
+5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections.
-6. **Cite Your Sources**: Track which documents/chunks informed your answer for citation.
+6. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function.
+7. **Cite Your Sources**: Track which documents/chunks informed your answer for citation.
 
 ## DoclingDocument API
 
@@ -114,6 +121,12 @@ summary = ask("What are the main machine learning approaches discussed?")
 print(summary)
 ```
 
+## Workflow
+
+1. **ALWAYS start by using execute_code** to explore the knowledge base
+2. Run multiple code blocks as needed to gather information
+3. After collecting data, provide your final answer
+
 ## Output Format
 
 After executing code and gathering information, provide:
@@ -121,4 +134,4 @@ After executing code and gathering information, provide:
 2. Key findings from your analysis
 3. References to specific documents/chunks that informed your answer
 
-Remember: You're solving problems that require computation, aggregation, or complex traversal - things traditional RAG can't do well. Write code to do the heavy lifting."""
+CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py

@@ -6,7 +6,8 @@ import traceback
 from io import StringIO
 from typing import TYPE_CHECKING, Any
 
-from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
+from haiku.rag.agents.rlm.dependencies import RLMContext
+from haiku.rag.config.models import RLMConfig
 
 if TYPE_CHECKING:
     from haiku.rag.client import HaikuRAG
@@ -151,11 +152,11 @@ class REPLEnvironment:
     def _make_search(self):
         """Create sync search function that bridges to async client."""
 
-        def search(
+        def search(query: str, limit: int = 10) -> list[dict]:
-            query: str, limit: int = 10, filter: str | None = None
-        ) -> list[dict]:
             async def _search():
-                return await self.client.search(query, limit=limit, filter=filter)
+                return await self.client.search(
+                    query, limit=limit, filter=self.context.filter
+                )
 
             results = self._run_async_from_thread(_search())
             self.context.search_results.extend(results)
@@ -178,12 +179,10 @@ class REPLEnvironment:
     def _make_list_documents(self):
         """Create sync list_documents function."""
 
-        def list_documents(
+        def list_documents(limit: int = 10, offset: int = 0) -> list[dict]:
-            limit: int = 10, offset: int = 0, filter: str | None = None
-        ) -> list[dict]:
             async def _list():
                 return await self.client.list_documents(
-                    limit=limit, offset=offset, filter=filter
+                    limit=limit, offset=offset, filter=self.context.filter
                 )
 
             docs = self._run_async_from_thread(_list())
@@ -250,9 +249,11 @@ class REPLEnvironment:
     def _make_ask(self):
         """Create sync ask function that uses QA agent."""
 
-        def ask(question: str, filter: str | None = None) -> str:
+        def ask(question: str) -> str:
             async def _ask():
-                answer, citations = await self.client.ask(question, filter=filter)
+                answer, citations = await self.client.ask(
+                    question, filter=self.context.filter
+                )
                 for c in citations:
                     for sr in self.context.search_results:
                         if sr.chunk_id == c.chunk_id:

haiku_rag_slim/haiku/rag/app.py

@@ -432,6 +432,37 @@ class HaikuRAGApp:
                     for renderable in format_citations_rich(citations):
                         self.console.print(renderable)
 
+    async def rlm(
+        self,
+        question: str,
+        document: str | None = None,
+        filter: str | None = None,
+    ):
+        """Answer a question using the RLM agent with code execution.
+
+        Args:
+            question: The question to answer
+            document: Optional document ID or title to pre-load
+            filter: SQL WHERE clause to filter documents
+        """
+        async with HaikuRAG(
+            db_path=self.db_path,
+            config=self.config,
+            read_only=self.read_only,
+            before=self.before,
+        ) as self.client:
+            documents = [document] if document else None
+
+            self.console.print(f"[bold blue]Question:[/bold blue] {question}")
+            self.console.print()
+            self.console.print("[dim]Running RLM agent with code execution...[/dim]")
+            self.console.print()
+
+            answer = await self.client.rlm(question, documents=documents, filter=filter)
+
+            self.console.print("[bold green]Answer:[/bold green]")
+            self.console.print(Markdown(answer))
+
     async def research(
         self,
         question: str,

haiku_rag_slim/haiku/rag/cli.py

@@ -364,6 +364,39 @@ def ask(
     )
 
 
+@_cli.command("rlm", help="Answer questions using code execution (RLM agent)")
+def rlm(
+    question: str = typer.Argument(
+        help="The question to answer",
+    ),
+    db: Path | None = typer.Option(
+        None,
+        "--db",
+        help="Path to the LanceDB database file",
+    ),
+    document: str | None = typer.Option(
+        None,
+        "--document",
+        "-d",
+        help="Document ID or title to pre-load for analysis",
+    ),
+    filter: str | None = typer.Option(
+        None,
+        "--filter",
+        "-f",
+        help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
+    ),
+):
+    app = create_app(db)
+    asyncio.run(
+        app.rlm(
+            question=question,
+            document=document,
+            filter=filter,
+        )
+    )
+
+
 @_cli.command("research", help="Run multi-agent research and output a concise report")
 def research(
     question: str = typer.Argument(..., help="The research question to investigate"),

haiku_rag_slim/haiku/rag/client.py

@@ -1293,6 +1293,53 @@ class HaikuRAG:
         qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
         return await qa_agent.answer(question, filter=filter)
 
+    async def rlm(
+        self,
+        question: str,
+        documents: list[str] | None = None,
+        filter: str | None = None,
+    ) -> str:
+        """Answer a question using the RLM agent with code execution.
+
+        The RLM (Recursive Language Model) agent can write and execute Python
+        code in a sandboxed environment to solve problems that require
+        computation, aggregation, or complex traversal across documents.
+
+        Args:
+            question: The question to answer.
+            documents: Optional list of document IDs or titles to pre-load.
+            filter: SQL WHERE clause to filter documents during searches.
+
+        Returns:
+            The answer as a string.
+        """
+        from haiku.rag.agents.rlm import RLMContext, RLMDeps, create_rlm_agent
+
+        context = RLMContext(filter=filter)
+
+        if documents:
+            loaded_docs = []
+            for doc_ref in documents:
+                doc = await self.get_document_by_id(doc_ref)
+                if not doc:
+                    docs = await self.list_documents(filter=f"title = '{doc_ref}'")
+                    if docs and docs[0].id:
+                        doc = await self.get_document_by_id(docs[0].id)
+                if doc:
+                    loaded_docs.append(doc)
+            context.documents = loaded_docs if loaded_docs else None
+
+        deps = RLMDeps(
+            client=self,
+            config=self._config,
+            context=context,
+        )
+
+        agent = create_rlm_agent(self._config)
+        result = await agent.run(question, deps=deps)
+
+        return result.output.answer
+
     async def visualize_chunk(self, chunk: Chunk) -> list:
         """Render page images with bounding box highlights for a chunk.
 

haiku_rag_slim/haiku/rag/config/models.py

@@ -94,6 +94,19 @@ class ResearchConfig(BaseModel):
     max_concurrency: int = 1
 
 
+class RLMConfig(BaseModel):
+    model: ModelConfig = Field(
+        default_factory=lambda: ModelConfig(
+            provider="ollama",
+            name="gpt-oss",
+            enable_thinking=False,
+        )
+    )
+    code_timeout: float = 60.0
+    max_output_chars: int = 50_000
+    max_tool_calls: int = 20
+
+
 class PictureDescriptionConfig(BaseModel):
     """Configuration for VLM-based picture description."""
 
@@ -194,6 +207,7 @@ class AppConfig(BaseModel):
     reranking: RerankingConfig = Field(default_factory=RerankingConfig)
     qa: QAConfig = Field(default_factory=QAConfig)
     research: ResearchConfig = Field(default_factory=ResearchConfig)
+    rlm: RLMConfig = Field(default_factory=RLMConfig)
     processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
     search: SearchConfig = Field(default_factory=SearchConfig)
     providers: ProvidersConfig = Field(default_factory=ProvidersConfig)

haiku_rag_slim/haiku/rag/mcp.py

@@ -245,4 +245,31 @@ def create_mcp_server(
         except Exception:
             return None
 
+    @mcp.tool()
+    async def rlm_question(
+        question: str,
+        document: str | None = None,
+        filter: str | None = None,
+    ) -> str:
+        """Answer complex questions using code execution (RLM agent).
+
+        Use this for questions requiring computation, aggregation, or
+        complex traversal across documents. The agent can write Python
+        code to search, analyze, and compute answers.
+
+        Args:
+            question: The question to answer.
+            document: Optional document ID or title to pre-load for analysis.
+            filter: Optional SQL WHERE clause to filter documents.
+
+        Returns:
+            The answer as a string.
+        """
+        try:
+            async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
+                documents = [document] if document else None
+                return await rag.rlm(question, documents=documents, filter=filter)
+        except Exception as e:
+            return f"Error running RLM agent: {e!s}"
+
     return mcp

tests/agents/rlm/conftest.py

@@ -1,8 +1,9 @@
 import pytest
 
-from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
+from haiku.rag.agents.rlm.dependencies import RLMContext
 from haiku.rag.agents.rlm.sandbox import REPLEnvironment
 from haiku.rag.client import HaikuRAG
+from haiku.rag.config.models import RLMConfig
 
 
 @pytest.fixture

tests/agents/rlm/test_agent.py

@@ -2,7 +2,7 @@ import pytest
 from pydantic_ai import Agent
 
 from haiku.rag.agents.rlm.agent import create_rlm_agent
-from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext, RLMDeps
+from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
 from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
 from haiku.rag.config import Config
 
@@ -26,12 +26,10 @@ class TestExecuteCodeTool:
         """Test that execute_code tool produces structured CodeExecution output."""
         from haiku.rag.agents.rlm.agent import _get_or_create_repl
 
-        config = RLMConfig()
         context = RLMContext()
         deps = RLMDeps(
             client=empty_client,
             config=Config,
-            rlm_config=config,
             context=context,
         )
 
@@ -51,12 +49,10 @@ class TestExecuteCodeTool:
         """Test that code executions are tracked as CodeExecution objects in RLMContext."""
         from haiku.rag.agents.rlm.agent import _get_or_create_repl
 
-        config = RLMConfig()
         context = RLMContext()
         deps = RLMDeps(
             client=empty_client,
             config=Config,
-            rlm_config=config,
             context=context,
         )
 
@@ -91,12 +87,10 @@ class TestExecuteCodeTool:
         """Test that failed executions are properly captured."""
         from haiku.rag.agents.rlm.agent import _get_or_create_repl
 
-        config = RLMConfig()
         context = RLMContext()
         deps = RLMDeps(
             client=empty_client,
             config=Config,
-            rlm_config=config,
             context=context,
         )
 

tests/agents/rlm/test_models.py

@@ -29,7 +29,6 @@ class TestRLMResult:
     def test_create_result_with_answer_only(self):
         result = RLMResult(answer="The answer is 42")
         assert result.answer == "The answer is 42"
-        assert result.citations == []
         assert result.code_executions == []
 
     def test_create_result_with_code_executions(self):
@@ -53,22 +52,3 @@ class TestRLMResult:
         )
         assert len(result.code_executions) == 2
         assert result.code_executions[1].stdout == "2\n"
-
-    def test_create_result_with_citations(self):
-        from haiku.rag.agents.research.models import Citation
-
-        citations = [
-            Citation(
-                document_id="doc1",
-                chunk_id="chunk1",
-                document_uri="file://test.pdf",
-                document_title="Test Doc",
-                content="Some content",
-            )
-        ]
-        result = RLMResult(
-            answer="Found in Test Doc",
-            citations=citations,
-        )
-        assert len(result.citations) == 1
-        assert result.citations[0].document_title == "Test Doc"

tests/agents/rlm/test_sandbox.py

@@ -445,6 +445,73 @@ class TestSandboxExecution:
         )  # Allow some margin for truncation message
 
 
+class TestContextFilter:
+    """Test that context filter is applied to all searches."""
+
+    @pytest.mark.asyncio
+    async def test_context_filter_applied_to_search(self, temp_db_path):
+        """Search applies context filter automatically."""
+        from unittest.mock import AsyncMock
+
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            context = RLMContext(filter="uri LIKE '%medical%'")
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+            client.search = AsyncMock(return_value=[])
+
+            await repl.execute_async("search('test query')")
+
+            client.search.assert_called_once_with(
+                "test query", limit=10, filter="uri LIKE '%medical%'"
+            )
+
+    @pytest.mark.asyncio
+    async def test_context_filter_applied_to_list_documents(self, temp_db_path):
+        """list_documents applies context filter automatically."""
+        from unittest.mock import AsyncMock
+
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            context = RLMContext(filter="title = 'Report'")
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+            client.list_documents = AsyncMock(return_value=[])
+
+            await repl.execute_async("list_documents()")
+
+            client.list_documents.assert_called_once_with(
+                limit=10, offset=0, filter="title = 'Report'"
+            )
+
+    @pytest.mark.asyncio
+    async def test_context_filter_applied_to_ask(self, temp_db_path):
+        """ask applies context filter automatically."""
+        from unittest.mock import AsyncMock
+
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            context = RLMContext(filter="metadata->>'category' = 'finance'")
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+            client.ask = AsyncMock(return_value=("Answer", []))
+
+            await repl.execute_async("ask('What is the revenue?')")
+
+            client.ask.assert_called_once_with(
+                "What is the revenue?", filter="metadata->>'category' = 'finance'"
+            )
+
+
 class TestSecurityEscapes:
     """Test that common security escape attempts are blocked."""