From 4118533db195aab072667b1065befe4ab00468d3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 16 Apr 2026 13:59:56 +0300 Subject: [PATCH] fold context expansion into sandbox search and remove get_context --- CHANGELOG.md | 4 ++- docs/agents/analysis.md | 3 +- .../haiku/rag/agents/analysis/agent.py | 2 +- .../haiku/rag/agents/analysis/prompts.py | 28 +++++++------------ .../haiku/rag/agents/analysis/sandbox.py | 15 ++-------- tests/agents/analysis/test_sandbox.py | 25 +++++++---------- ...test_search_returns_expanded_content.yaml} | 0 7 files changed, 27 insertions(+), 50 deletions(-) rename tests/cassettes/test_sandbox/{TestSandboxGetContext.test_get_context_returns_expanded_content.yaml => TestSandboxSearchExpandsContext.test_search_returns_expanded_content.yaml} (100%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 583c2faf..2fc6327c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ ### Changed +- **Analysis sandbox `search()` now returns expanded results**: Search results automatically include surrounding context (adjacent paragraphs, complete tables, section content) via the document_items table + - **BREAKING**: Rename RLM agent to analysis agent throughout: - `agents/rlm/` → `agents/analysis/`, all classes renamed (`RLMResult` → `AnalysisResult`, etc.) - `client.rlm()` → `client.analyze()` @@ -14,7 +16,7 @@ ### Removed -- **`get_chunk()`**: Removed from analysis sandbox +- **`get_chunk()`**: Removed from analysis sandbox — search results now include expanded context automatically - **`create_analysis_toolset()`**: Removed unused `tools/analysis.py` module. ## [0.40.1] - 2026-04-17 diff --git a/docs/agents/analysis.md b/docs/agents/analysis.md index 455eb7f9..333e44d0 100644 --- a/docs/agents/analysis.md +++ b/docs/agents/analysis.md @@ -58,8 +58,7 @@ The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https: | Function | Description | |----------|-------------| -| `search(query, limit)` | Hybrid search (vector + full-text) returning matching chunks with scores | -| `get_context(chunk_id)` | Expand a chunk with surrounding content (adjacent paragraphs, complete tables) | +| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion | | `list_documents(limit, offset)` | List documents in the knowledge base | | `get_document(id_or_title)` | Get full text content of a document | | `get_docling_document(document_id)` | Get the DoclingDocument structure as a dict (texts, tables, pictures) | diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/agent.py b/haiku_rag_slim/haiku/rag/agents/analysis/agent.py index 66d11b4d..2662fca5 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/agent.py @@ -34,7 +34,7 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, AnalysisResu async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution: """Execute Python code in a sandboxed interpreter. - The code has access to haiku.rag functions (search, get_context, + The code has access to haiku.rag functions (search, list_documents, get_document, get_docling_document, llm). Use print() to output results. diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py index cee51c01..00ca88cd 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py @@ -11,12 +11,9 @@ Inside execute_code, these functions are ALREADY available in the namespace. Do ### await search(query, limit=10) -> list[dict] Search the knowledge base using hybrid search (vector + full-text). +Results are automatically expanded with surrounding context (adjacent paragraphs, complete tables, section content). Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings -### await get_context(chunk_id) -> str | None -Get expanded content around a chunk, including surrounding paragraphs, complete tables, and adjacent sections from the same document. -Use this after search() when a result looks relevant but you need more context to understand it fully. - ### await 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 @@ -57,26 +54,21 @@ For pattern matching or text extraction, use `import re`, string methods (`str.s ## Strategy Guide -1. **Search First**: Start with `search()` to find relevant content. Examine the results to understand what's available. -2. **Expand When Needed**: If a search result looks relevant but incomplete, use `get_context(chunk_id)` to get surrounding content from the same document. -3. **Use get_document for Full Text**: When you need a document's complete text (e.g., for regex across the whole document), use `get_document(id_or_title)`. -4. **Use get_docling_document for Structure**: When you need structured data like table grids, document hierarchy, or section labels, use `get_docling_document(document_id)`. -5. **Iterate**: Run code, examine results, refine your approach. Don't try to solve everything in one execution. -6. **Use llm() for Reasoning**: When you have content and need classification, summarization, or extraction, use `llm()` rather than writing complex parsing logic. -7. **Document Titles Are Often None**: Use `uri` or `id` to identify documents. Use `list_documents()` to discover what's available. +1. **Search First**: Start with `search()` to find relevant content. Results already include expanded context (surrounding paragraphs, complete tables, section content). +2. **Use get_document for Full Text**: When you need a document's complete text (e.g., for regex across the whole document), use `get_document(id_or_title)`. +3. **Use get_docling_document for Structure**: When you need structured data like table grids, document hierarchy, or section labels, use `get_docling_document(document_id)`. +4. **Iterate**: Run code, examine results, refine your approach. Don't try to solve everything in one execution. +5. **Use llm() for Reasoning**: When you have content and need classification, summarization, or extraction, use `llm()` rather than writing complex parsing logic. +6. **Document Titles Are Often None**: Use `uri` or `id` to identify documents. Use `list_documents()` to discover what's available. ## Example Patterns -### Search and expand context +### Search (results include expanded context) ```python results = await search("revenue figures", limit=5) for r in results: - print(f"{r['document_title']}: {r['content'][:100]}") - -# Get more context around the most relevant result -expanded = await get_context(results[0]['chunk_id']) -if expanded: - print(f"Expanded: {expanded[:500]}") + print(f"{r['document_title']} (score={r['score']:.2f}):") + print(r['content'][:200]) ``` ### Extracting data with regex diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index e0f43547..3f5ca223 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -7,7 +7,6 @@ import pydantic_monty from haiku.rag.agents.analysis.dependencies import AnalysisContext from haiku.rag.config.models import AppConfig from haiku.rag.store.compression import decompress_json -from haiku.rag.store.models.chunk import SearchResult if TYPE_CHECKING: from haiku.rag.client import HaikuRAG @@ -55,6 +54,7 @@ class Sandbox: async def search(query: str, limit: int = 10) -> list[dict[str, Any]]: results = await client.search(query, limit=limit, filter=context.filter) + expanded = await client.expand_context(results) return [ { "chunk_id": r.chunk_id, @@ -66,7 +66,7 @@ class Sandbox: "page_numbers": r.page_numbers, "headings": r.headings, } - for r in results + for r in expanded ] async def list_documents( @@ -89,16 +89,6 @@ class Sandbox: doc = await client.resolve_document(id_or_title) return doc.content if doc else None - async def get_context(chunk_id: str) -> str | None: - chunk = await client.get_chunk_by_id(chunk_id) - if not chunk: - return None - search_result = SearchResult.from_chunk(chunk, score=1.0) - expanded = await client.expand_context([search_result]) - if expanded: - return expanded[0].content - return chunk.content - async def get_docling_document( document_id: str, ) -> dict[str, Any] | None: @@ -122,7 +112,6 @@ class Sandbox: "search": search, "list_documents": list_documents, "get_document": get_document, - "get_context": get_context, "get_docling_document": get_docling_document, "llm": llm, } diff --git a/tests/agents/analysis/test_sandbox.py b/tests/agents/analysis/test_sandbox.py index f40b0f6e..2e975f1f 100644 --- a/tests/agents/analysis/test_sandbox.py +++ b/tests/agents/analysis/test_sandbox.py @@ -163,22 +163,19 @@ class TestSandboxHaikuRAG: assert "True" in result.stdout -class TestSandboxGetContext: - """Test get_context() external function.""" +class TestSandboxSearchExpandsContext: + """Test that search() returns expanded results.""" @pytest.mark.asyncio - async def test_get_context_missing_chunk(self, sandbox): - """get_context returns None for a non-existent chunk.""" - result = await sandbox.execute( - "ctx = await get_context('nonexistent-id')\nprint(ctx is None)" - ) - assert result.success - assert "True" in result.stdout + async def test_get_context_not_available(self, sandbox): + """get_context is no longer a sandbox function.""" + result = await sandbox.execute("await get_context('x')") + assert not result.success @pytest.mark.asyncio @pytest.mark.vcr() - async def test_get_context_returns_expanded_content(self, temp_db_path): - """get_context returns content for a valid chunk.""" + async def test_search_returns_expanded_content(self, temp_db_path): + """search() returns context-expanded results.""" config = AppConfig() async with HaikuRAG(temp_db_path, create=True) as client: await client.create_document( @@ -191,10 +188,8 @@ class TestSandboxGetContext: sb = Sandbox(client=client, config=config, context=context) result = await sb.execute( "results = await search('fox', limit=1)\n" - "chunk_id = results[0]['chunk_id']\n" - "ctx = await get_context(chunk_id)\n" - "print(type(ctx).__name__)\n" - "print('fox' in ctx.lower())" + "print(type(results[0]['content']).__name__)\n" + "print('fox' in results[0]['content'].lower())" ) assert result.success assert "str" in result.stdout diff --git a/tests/cassettes/test_sandbox/TestSandboxGetContext.test_get_context_returns_expanded_content.yaml b/tests/cassettes/test_sandbox/TestSandboxSearchExpandsContext.test_search_returns_expanded_content.yaml similarity index 100% rename from tests/cassettes/test_sandbox/TestSandboxGetContext.test_get_context_returns_expanded_content.yaml rename to tests/cassettes/test_sandbox/TestSandboxSearchExpandsContext.test_search_returns_expanded_content.yaml