fold context expansion into sandbox search and remove get_context

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

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) |

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.

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

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,
         }

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