diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/__init__.py b/haiku_rag_slim/haiku/rag/agents/rlm/__init__.py index 1d1e5038..77ba8c73 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/__init__.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/agent.py b/haiku_rag_slim/haiku/rag/agents/rlm/agent.py index ccff6c85..9fc55309 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/agent.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py b/haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py index c5dda1a8..8f022f32 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/models.py b/haiku_rag_slim/haiku/rag/agents/rlm/models.py index d6e31851..f3f1793e 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/models.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py b/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py index 7f21d636..90171fbd 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py +++ b/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. -2. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find. -3. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with. -4. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections. -5. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function. -6. **Cite Your Sources**: Track which documents/chunks informed your answer for citation. +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. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content. +3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find. +4. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with. +5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections. +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.""" diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py index 31f19d44..5daf7e2a 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py +++ b/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( - query: str, limit: int = 10, filter: str | None = None - ) -> list[dict]: + def search(query: str, limit: int = 10) -> 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( - limit: int = 10, offset: int = 0, filter: str | None = None - ) -> list[dict]: + def list_documents(limit: int = 10, offset: int = 0) -> 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: diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index b627b694..067cfddf 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 7b9cdd27..8e688d32 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/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"), diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index da0ffa77..6d39c62e 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index e5e56b2f..7d2c1828 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index b0d9fcf9..9bd1b34b 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/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 diff --git a/tests/agents/rlm/conftest.py b/tests/agents/rlm/conftest.py index 580671c5..7efb0a86 100644 --- a/tests/agents/rlm/conftest.py +++ b/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 diff --git a/tests/agents/rlm/test_agent.py b/tests/agents/rlm/test_agent.py index 1e0a7190..2816ca6a 100644 --- a/tests/agents/rlm/test_agent.py +++ b/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, ) diff --git a/tests/agents/rlm/test_models.py b/tests/agents/rlm/test_models.py index 73a11213..46767473 100644 --- a/tests/agents/rlm/test_models.py +++ b/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" diff --git a/tests/agents/rlm/test_sandbox.py b/tests/agents/rlm/test_sandbox.py index c158a80e..bacb326b 100644 --- a/tests/agents/rlm/test_sandbox.py +++ b/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."""