Add get_docling_document external function to sandbox

CHANGELOG.md

@@ -4,7 +4,7 @@
 ### Changed
 
 - **RLM sandbox**: Replaced Docker-based code execution with [pydantic-monty](https://github.com/pydantic/monty), a minimal secure Python interpreter written in Rust. Eliminates Docker as a runtime dependency for RLM with sub-millisecond sandbox startup
-- **RLM sandbox functions**: Replaced `get_docling_document()` with `get_chunk(chunk_id)` for retrieving chunk content and metadata from search results
+- **RLM sandbox functions**: Added `get_chunk(chunk_id)` for retrieving chunk content and metadata from search results. `get_docling_document(document_id)` now returns the full document structure as a JSON dict
 - **`RLMConfig`**: Removed `docker_image` and `docker_memory_limit` fields
 
 ### Added

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

@@ -28,6 +28,16 @@ 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.
+Top-level keys: name, texts, tables, pictures, body, pages, key_value_items, furniture, groups.
+- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+- `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
+Use this for structural analysis: extracting table data, counting sections, analyzing layout.
+
 ### await llm(prompt) -> str
 Call an LLM directly with the given prompt. Returns the response as a string.
 Use this for classification, summarization, extraction, or any task where you
@@ -97,6 +107,22 @@ for r in results:
     print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
 ```
 
+### Extracting tables from a document
+```python
+docs = await list_documents(limit=10)
+for d in docs:
+    doc = await get_docling_document(d['id'])
+    if doc:
+        tables = doc.get('tables', [])
+        if tables:
+            print(f"{d['title']}: {len(tables)} table(s)")
+            for i, table in enumerate(tables):
+                grid = table.get('data', {}).get('grid', [])
+                for row in grid:
+                    cells = [cell.get('text', '') for cell in row]
+                    print(f"  Table {i}: {cells}")
+```
+
 ### Using llm() for classification
 ```python
 content = await get_document("Q1 Report")

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

@@ -1,3 +1,4 @@
+import json
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal
 
@@ -5,6 +6,7 @@ import pydantic_monty
 
 from haiku.rag.agents.rlm.dependencies import RLMContext
 from haiku.rag.config.models import AppConfig
+from haiku.rag.store.compression import decompress_json
 
 if TYPE_CHECKING:
     from haiku.rag.client import HaikuRAG
@@ -106,6 +108,15 @@ class Sandbox:
                 "labels": meta.labels,
             }
 
+        async def get_docling_document(
+            document_id: str,
+        ) -> dict[str, Any] | None:
+            doc = await client.get_document_by_id(document_id)
+            if not doc or not doc.docling_document:
+                return None
+            json_str = decompress_json(doc.docling_document)
+            return json.loads(json_str)
+
         async def llm(prompt: str) -> str:
             from pydantic_ai import Agent
 
@@ -121,6 +132,7 @@ class Sandbox:
             "list_documents": list_documents,
             "get_document": get_document,
             "get_chunk": get_chunk,
+            "get_docling_document": get_docling_document,
             "llm": llm,
         }
 

tests/agents/rlm/test_sandbox.py

@@ -337,6 +337,43 @@ class TestSandboxPreloadedDocuments:
         assert "Doc B" in result.stdout
 
 
+class TestSandboxDoclingDocument:
+    """Test get_docling_document() external function."""
+
+    @pytest.mark.asyncio
+    async def test_returns_none_for_missing_document(self, sandbox):
+        """get_docling_document returns None for a non-existent document."""
+        result = await sandbox.execute(
+            "doc = await get_docling_document('nonexistent-id')\nprint(doc is None)"
+        )
+        assert result.success
+        assert "True" in result.stdout
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_returns_dict_for_document_with_docling_data(self, temp_db_path):
+        """get_docling_document returns a dict for a document with docling data."""
+        config = AppConfig()
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            doc = await client.create_document(
+                content="Docling processed content",
+                uri="test://docling",
+                title="Docling Doc",
+            )
+
+            context = RLMContext()
+            sb = Sandbox(client=client, config=config, context=context)
+            result = await sb.execute(
+                f"doc = await get_docling_document('{doc.id}')\n"
+                "print(type(doc).__name__)\n"
+                "print(doc['name'])\n"
+                "print('texts' in doc)"
+            )
+            assert result.success
+            assert "dict" in result.stdout
+            assert "True" in result.stdout
+
+
 class TestSandboxLLM:
     """Test llm() external function."""