add get_context() to analysis sandbox and improve prompt

CHANGELOG.md

@@ -1,6 +1,22 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- **BREAKING**: Rename RLM agent to analysis agent throughout:
+  - `agents/rlm/` → `agents/analysis/`, all classes renamed (`RLMResult` → `AnalysisResult`, etc.)
+  - `client.rlm()` → `client.analyze()`
+  - CLI: `haiku-rag rlm` → `haiku-rag analyze`
+  - MCP: `rlm_question` → `analyze`
+  - Config: `rlm:` → `analysis:` in YAML, `RLMConfig` → `AnalysisConfig`
+  - Skill: `rag-rlm` → `rag-analysis`, `skills/rlm.py` → `skills/analysis.py`
+  - State namespace: `"rlm"` → `"analysis"`
+
+### Removed
+
+- **`get_chunk()`**: Removed from analysis sandbox
+- **`create_analysis_toolset()`**: Removed unused `tools/analysis.py` module.
+
 ## [0.40.1] - 2026-04-17
 
 ### Fixed
@@ -21,21 +37,12 @@
 - **`max_searches` default**: Raised from 3 to 5 — faster expansion makes additional searches inexpensive
 - **Improved QA prompt**: Stronger instruction to refuse answering from tangentially related content
 - **Improved judge prompt**: Asymmetric evaluation — generated answers that are more comprehensive than expected are not penalized
-- **BREAKING**: Rename RLM agent to analysis agent throughout:
-  - `agents/rlm/` → `agents/analysis/`, all classes renamed (`RLMResult` → `AnalysisResult`, etc.)
-  - `client.rlm()` → `client.analyze()`
-  - CLI: `haiku-rag rlm` → `haiku-rag analyze`
-  - MCP: `rlm_question` → `analyze`
-  - Config: `rlm:` → `analysis:` in YAML, `RLMConfig` → `AnalysisConfig`
-  - Skill: `rag-rlm` → `rag-analysis`, `skills/rlm.py` → `skills/analysis.py`
-  - State namespace: `"rlm"` → `"analysis"`
 
 ### Removed
 
 - **`context_radius` config**: Replaced by automatic section-bounded expansion. Context expansion no longer requires configuration.
 - **DoclingDocument LRU cache**: No longer needed — the document_items table replaces in-memory caching for context expansion
 - **`cachetools` dependency**: No longer used
-- **`create_analysis_toolset()`**: Removed unused `tools/analysis.py` module.
 
 ## [0.39.0] - 2026-04-09
 

docs/agents/analysis.md

@@ -59,9 +59,9 @@ 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) |
 | `list_documents(limit, offset)` | List documents in the knowledge base |
 | `get_document(id_or_title)` | Get full text content of a document |
-| `get_chunk(chunk_id)` | Get a chunk with metadata (headings, page numbers, labels) for citations |
 | `get_docling_document(document_id)` | Get the DoclingDocument structure as a dict (texts, tables, pictures) |
 | `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
 

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

@@ -34,8 +34,8 @@ 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, list_documents,
-        get_document, get_chunk, llm).
+        The code has access to haiku.rag functions (search, get_context,
+        list_documents, get_document, get_docling_document, llm).
 
         Use print() to output results.
 

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

@@ -13,6 +13,10 @@ Inside execute_code, these functions are ALREADY available in the namespace. Do
 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
 
+### 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
@@ -21,18 +25,12 @@ Returns list of dicts with keys: id, title, uri, created_at
 Get the full text content of a document by ID, title, or URI.
 Returns the document content as a string, or None if not found.
 
-### await get_chunk(chunk_id) -> dict | None
-Get a specific chunk by its ID (from search results).
-Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
-Use this to retrieve full chunk details and metadata for citation.
-
 ### await get_docling_document(document_id) -> dict | None
 Get the full document structure as a dict (DoclingDocument format).
 Use `list_documents()` or search results to get document IDs first.
-- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item")
 - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
 - `pictures`: list of figures/images with metadata
-- `pages`: page dimensions and metadata
 
 ### await llm(prompt) -> str
 Call an LLM directly with the given prompt. Returns the response as a string.
@@ -59,24 +57,26 @@ For pattern matching or text extraction, use `import re`, string methods (`str.s
 
 ## Strategy Guide
 
-1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
-3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
-4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
+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.
 
 ## Example Patterns
 
-### Counting documents matching a condition
+### Search and expand context
 ```python
-docs = await list_documents(limit=100)
-count = 0
-for doc in docs:
-    content = await get_document(doc['id'])
-    if content and 'keyword' in content.lower():
-        count += 1
-        print(f"Found in: {doc['title']}")
-print(f"Total: {count}")
+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]}")
 ```
 
 ### Extracting data with regex
@@ -108,6 +108,15 @@ for d in docs:
                     print(f"  Table {i}: {cells}")
 ```
 
+### Regex search across a full document
+```python
+import re
+content = await get_document("Policy Document")
+if content:
+    emails = re.findall(r'[\\w.+-]+@[\\w-]+\\.[\\w.]+', content)
+    print(f"Found {len(emails)} email addresses: {emails}")
+```
+
 ## Output Format
 
 Your final response MUST be valid JSON matching this exact schema:

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

@@ -7,6 +7,7 @@ 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
@@ -88,25 +89,15 @@ class Sandbox:
             doc = await client.resolve_document(id_or_title)
             return doc.content if doc else None
 
-        async def get_chunk(chunk_id: str) -> dict[str, Any] | None:
+        async def get_context(chunk_id: str) -> str | None:
             chunk = await client.get_chunk_by_id(chunk_id)
             if not chunk:
                 return None
-            meta = chunk.get_chunk_metadata()
-            doc_title = chunk.document_title
-            if not doc_title and chunk.document_id:
-                doc = await client.get_document_by_id(chunk.document_id)
-                if doc:
-                    doc_title = doc.title
-            return {
-                "chunk_id": chunk.id,
-                "content": chunk.content,
-                "document_id": chunk.document_id,
-                "document_title": doc_title,
-                "headings": meta.headings,
-                "page_numbers": meta.page_numbers,
-                "labels": meta.labels,
-            }
+            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,
@@ -131,7 +122,7 @@ class Sandbox:
             "search": search,
             "list_documents": list_documents,
             "get_document": get_document,
-            "get_chunk": get_chunk,
+            "get_context": get_context,
             "get_docling_document": get_docling_document,
             "llm": llm,
         }

tests/agents/analysis/test_agent.py

@@ -139,17 +139,10 @@ class TestClientAnalysisIntegration:
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_analyze_search_and_get_chunk(
+    async def test_analyze_search_and_identify_source(
         self, allow_model_requests, temp_db_path
     ):
-        """Test analysis agent can search and use get_chunk for citations.
-
-        Agent program:
-            results = search("content", limit=5)
-            for r in results:
-                chunk = get_chunk(r['chunk_id'])
-                print(chunk['document_title'], chunk['chunk_id'])
-        """
+        """Test analysis agent can search and identify source documents."""
         from haiku.rag.client import HaikuRAG
 
         config = AppConfig()

tests/agents/analysis/test_sandbox.py

@@ -162,41 +162,44 @@ class TestSandboxHaikuRAG:
         assert result.success
         assert "True" in result.stdout
 
+
+class TestSandboxGetContext:
+    """Test get_context() external function."""
+
+    @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
+
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_get_chunk(self, temp_db_path):
-        """Test get_chunk function returns chunk with metadata."""
+    async def test_get_context_returns_expanded_content(self, temp_db_path):
+        """get_context returns content for a valid chunk."""
         config = AppConfig()
         async with HaikuRAG(temp_db_path, create=True) as client:
             await client.create_document(
-                content="Content about foxes and dogs.",
-                uri="test://doc",
-                title="Fox Document",
+                content="The quick brown fox jumps over the lazy dog.",
+                uri="test://animals",
+                title="Animals",
             )
 
             context = AnalysisContext()
             sb = Sandbox(client=client, config=config, context=context)
-            # First search to get a chunk_id
             result = await sb.execute(
-                "results = await search('foxes', limit=1)\n"
+                "results = await search('fox', limit=1)\n"
                 "chunk_id = results[0]['chunk_id']\n"
-                "chunk = await get_chunk(chunk_id)\n"
-                "print(chunk['document_title'])\n"
-                "print('content' in chunk)"
+                "ctx = await get_context(chunk_id)\n"
+                "print(type(ctx).__name__)\n"
+                "print('fox' in ctx.lower())"
             )
             assert result.success
-            assert "Fox Document" in result.stdout
+            assert "str" in result.stdout
             assert "True" in result.stdout
 
-    @pytest.mark.asyncio
-    async def test_get_chunk_not_found(self, sandbox):
-        """Test get_chunk returns None for missing chunk."""
-        result = await sandbox.execute(
-            "chunk = await get_chunk('nonexistent-id')\nprint(chunk is None)"
-        )
-        assert result.success
-        assert "True" in result.stdout
-
 
 class TestSandboxExternalFunctionEdgeCases:
     """Test edge cases in external function dispatch."""

tests/skills/test_analysis.py

@@ -155,7 +155,29 @@ class TestAnalyzeTool:
         assert state.analyses[0].answer == "42"
         assert state.analyses[0].program == "print(42)"
 
-    async def test_analyze_with_document_filter_in_state(self, rag_db, monkeypatch):
+    async def test_analyze_applies_document_filter_from_state(
+        self, rag_db, monkeypatch
+    ):
+        from haiku.rag.skills.analysis import AnalysisState, create_skill
+
+        captured_kwargs = {}
+
+        async def mock_analyze(self, question, **kwargs):
+            captured_kwargs.update(kwargs)
+            return AnalysisResult(answer="42", program="print(42)")
+
+        monkeypatch.setattr(HaikuRAG, "analyze", mock_analyze)
+
+        skill = create_skill(db_path=rag_db)
+        analyze = _get_tool(skill, "analyze")
+        state = AnalysisState(document_filter="title = 'AI Overview'")
+        ctx = _make_ctx(state)
+        await analyze(ctx, question="How many documents?")
+        assert captured_kwargs.get("filter") == "title = 'AI Overview'"
+
+    async def test_analyze_combines_state_filter_with_explicit_filter(
+        self, rag_db, monkeypatch
+    ):
         from haiku.rag.skills.analysis import AnalysisState, create_skill
 
         captured_kwargs = {}