danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add document virtual filesystem to analysis sandbox

Browse source 
Replace get_document() and get_docling_document() with a VFS at
/documents/{id}/ with metadata.json (eager), content.txt (lazy),
and items.jsonl (lazy). Keep search(), list_documents() (now returns
all), and llm() as external functions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yiorgis Gozadinos 2026-04-16 15:29:53 +03:00
parent bacc21b38b
commit a45820dbf7
No known key found for this signature in database
15 changed files with 527 additions and 203 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

@ -1,10 +1,14 @@
# Changelog
## [Unreleased]
### Added
- **Document virtual filesystem in analysis sandbox**: Documents are mounted at `/documents/{id}/` with `metadata.json` (eager), `content.txt` (lazy), and `items.jsonl` (lazy). The agent uses standard Python `pathlib.Path` to browse and read document content and structure.
- **`doc_item_refs` and `labels` in search results**: Search results now include document item references and labels for cross-referencing with `items.jsonl`.
### 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()`
@ -16,6 +20,7 @@
### Removed
- **`get_document()`, `get_docling_document()`**: Removed from analysis sandbox — replaced by the document virtual filesystem
- **`get_chunk()`**: Removed from analysis sandbox — search results now include expanded context automatically
- **`create_analysis_toolset()`**: Removed unused `tools/analysis.py` module.

41
docs/agents/analysis.md
View file

@ -11,7 +11,7 @@ The analysis agent enables complex analytical tasks by writing and executing Pyt
1. The agent receives a question
2. It writes Python code to explore the knowledge base
3. Code executes in a sandboxed Python interpreter with access to knowledge base functions
3. Code executes in a sandboxed Python interpreter with access to search, LLM, and a virtual filesystem of documents
4. The agent iterates: run code, examine results, refine approach
5. Final answer is synthesized from the gathered data
@ -54,37 +54,54 @@ async with HaikuRAG(path_to_db) as client:
## Sandbox Capabilities
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with access to these knowledge base functions:
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with:
### Functions
| Function | Description |
|----------|-------------|
| `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) |
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` for cross-referencing with `items.jsonl` |
| `list_documents()` | List all documents in the knowledge base |
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
When documents are pre-loaded via the `documents` parameter, they are injected as a `documents` variable accessible in the sandbox code.
### Document Filesystem
All documents are mounted as a virtual filesystem at `/documents/`. The agent uses standard Python `pathlib.Path` to browse and read files:
```
/documents/{document_id}/
metadata.json # {id, title, uri, created_at}
content.txt # Full document text
items.jsonl # Structured items: position, self_ref, label, text, page_numbers
```
- **`metadata.json`** — Loaded eagerly (small). Use `Path('/documents').iterdir()` to discover documents.
- **`content.txt`** — Lazy-loaded on first read. Full document text for regex or keyword search.
- **`items.jsonl`** — Lazy-loaded on first read. One JSON object per line with structured document elements. Tables are pre-rendered as markdown. Labels include `section_header`, `text`, `table`, `list_item`, `caption`, `formula`, `picture`, `code`, `footnote`, etc.
Search results include `doc_item_refs` (e.g. `["#/texts/5", "#/tables/0"]`) that match `self_ref` values in `items.jsonl`, enabling navigation from search hits to document structure.
When documents are pre-loaded via the `documents` parameter, they are also injected as a `documents` variable accessible in the sandbox code.
### Python Features
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, and the `json`, `re`, `math` modules.
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, file I/O via `pathlib.Path`, and the `json`, `re`, `math` modules.
Not supported: most imports (only `json`, `re`, `math` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
### Security
Code executes in an isolated interpreter with:
- **No filesystem access**: Code cannot read or write files
- **Virtual filesystem only**: The `/documents/` filesystem is sandboxed — no access to the real filesystem
- **No network access**: Code cannot make HTTP requests or open sockets
- **No imports**: Only `json`, `re`, and `math` modules are available
- **No imports**: Only `json`, `re`, `math`, and `pathlib` modules are available
- **Execution timeout**: Configurable limit (default 60s)
- **Output truncation**: Large outputs are truncated to prevent memory issues
## Context Filter
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM:
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
```python
# Agent can only see documents with "confidential" in the URI

4
haiku_rag_slim/haiku/rag/agents/analysis/agent.py
View file

@ -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_docling_document, llm).
The code has access to search() and llm() functions, and a
virtual filesystem at /documents/ with document content and structure.
Use print() to output results.

153
haiku_rag_slim/haiku/rag/agents/analysis/prompts.py
View file

@ -1,39 +1,92 @@
ANALYSIS_SYSTEM_PROMPT = """You are an analysis agent that solves complex research questions by writing and executing Python code.
You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
You MUST use the `execute_code` tool to run Python code. The functions and filesystem described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
## Available Functions
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
- results = await search("query") ✓ CORRECT
- import search ✗ WRONG - will fail
- results = search("query") ✗ WRONG - must use await
## Available Functions
### 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
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
### await list_documents(limit=10, offset=0) -> list[dict]
List available documents in the knowledge base.
### await list_documents() -> list[dict]
List all documents in the knowledge base.
Returns list of dicts with keys: id, title, uri, created_at
### await get_document(id_or_title) -> str | None
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_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")
- `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
### 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
already have the content and just need LLM reasoning.
## Document Filesystem
All documents in the knowledge base are available as files under `/documents/`. Use `from pathlib import Path` and standard file I/O to access them.
### Directory structure
```
/documents/
{document_id}/
metadata.json # {"id", "title", "uri", "created_at"}
content.txt # Full document text
items.jsonl # Structured document items (one JSON object per line)
```
### metadata.json
Small file with document metadata. Use to discover and identify documents.
```python
from pathlib import Path
import json
for doc_dir in Path('/documents').iterdir():
meta = json.loads((doc_dir / 'metadata.json').read_text())
print(meta['title'], meta['uri'])
```
### content.txt
Full text content of the document. Use for regex, keyword search, or full-text analysis.
```python
content = Path(f'/documents/{doc_id}/content.txt').read_text()
```
### items.jsonl
Structured document items as JSONL. Each line is a JSON object with:
- `position`: sequential position in the document
- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote", etc.
- `text`: rendered content (tables are markdown with `|` columns)
- `page_numbers`: list of page numbers where the item appears
Use items.jsonl to find tables, section headers, or specific structural elements:
```python
import json
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
for line in items_text.strip().split(chr(10)):
item = json.loads(line)
if item['label'] == 'table':
print(f"Table on page {item['page_numbers']}: {item['text'][:100]}")
```
## Cross-referencing search results with items
Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. Use this to navigate from a search hit to the surrounding document structure:
```python
results = await search("revenue", limit=5)
r = results[0]
doc_id = r['document_id']
refs = set(r['doc_item_refs'])
import json
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
for line in items_text.strip().split(chr(10)):
item = json.loads(line)
if item['self_ref'] in refs:
print(f"Matched: {item['label']} on page {item['page_numbers']}")
```
## Pre-loaded Documents Variable
If documents were pre-loaded for this session, a `documents` variable is available:
@ -46,68 +99,18 @@ Check if it exists with: `try: documents ... except NameError: ...`
## Available Python Features
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `filter()`, `getattr()`, `sorted()`/`.sort(key=...)`, try/except, and the `json`, `re`, `math` modules.
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `filter()`, `getattr()`, `sorted()`/`.sort(key=...)`, try/except, and the `json`, `re`, `math` modules. File I/O via `pathlib.Path` is supported for the `/documents/` filesystem.
Not supported: most imports (only `json`, `re`, `math` are available), class definitions, generators/yield, match statements, decorators, `with` statements.
For pattern matching or text extraction, use `import re`, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements.
## Strategy Guide
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 (results include expanded context)
```python
results = await search("revenue figures", limit=5)
for r in results:
print(f"{r['document_title']} (score={r['score']:.2f}):")
print(r['content'][:200])
```
### Extracting data with regex
```python
import re
numbers = []
results = await search("financial data", limit=20)
for r in results:
amounts = re.findall(r'\\$([\\d,]+)', r['content'])
for a in amounts:
numbers.append(int(a.replace(',', '')))
if numbers:
print(f"Average: {sum(numbers) / len(numbers)}")
```
### 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}")
```
### 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}")
```
1. **Search First**: Start with `search()` to find relevant content. Results include expanded context and `doc_item_refs` for cross-referencing.
2. **Discover Documents**: Use `list_documents()` to see what's in the knowledge base.
3. **Use items.jsonl for Structure**: Find tables, section headers, or specific elements by label and page number. Tables are pre-rendered as markdown.
4. **Use content.txt for Full Text**: When you need the complete document text (e.g., for regex across the whole document).
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.
## Output Format

141
haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py
View file

@ -1,14 +1,19 @@
import asyncio
import concurrent.futures
import json
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from pydantic_monty import CallbackFile, MemoryFile, OSAccess
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.config.models import AppConfig
from haiku.rag.store.compression import decompress_json
if TYPE_CHECKING:
from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@ -21,12 +26,19 @@ class SandboxResult:
success: bool
def _run_async(coro: Any) -> Any:
"""Run an async coroutine from a sync context (CallbackFile read)."""
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
return pool.submit(asyncio.run, coro).result()
class Sandbox:
"""Execute code in a sandboxed Python interpreter.
Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
External functions (search, list_documents, etc.) are called by Monty code
using ``await`` and resolved asynchronously on the host.
External functions (search, llm) are called by Monty code using ``await``
and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
sandbox = Sandbox(client, config, context)
result = await sandbox.execute("print('hello')")
@ -71,12 +83,8 @@ class Sandbox:
for r in expanded
]
async def list_documents(
limit: int = 10, offset: int = 0
) -> list[dict[str, Any]]:
docs = await client.list_documents(
limit=limit, offset=offset, filter=context.filter
)
async def list_documents() -> list[dict[str, Any]]:
docs = await client.list_documents(filter=context.filter)
return [
{
"id": d.id,
@ -87,19 +95,6 @@ class Sandbox:
for d in docs
]
async def get_document(id_or_title: str) -> str | None:
doc = await client.resolve_document(id_or_title)
return doc.content if doc else None
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
@ -113,14 +108,111 @@ class Sandbox:
return {
"search": search,
"list_documents": list_documents,
"get_document": get_document,
"get_docling_document": get_docling_document,
"llm": llm,
}
async def _build_vfs(self) -> OSAccess:
"""Build the virtual filesystem with document data.
Mounts per-document directories with:
- metadata.json: MemoryFile (eager, small)
- content.txt: CallbackFile (lazy, can be large)
- items.jsonl: CallbackFile (lazy, can be large)
"""
client = self._client
files: list[MemoryFile | CallbackFile] = []
docs = await client.list_documents(filter=self._context.filter)
for doc in docs:
if not doc.id:
continue
doc_id: str = doc.id
doc_dir = f"/documents/{doc_id}"
metadata = json.dumps(
{
"id": doc_id,
"title": doc.title,
"uri": doc.uri,
"created_at": str(doc.created_at),
},
ensure_ascii=False,
)
files.append(MemoryFile(f"{doc_dir}/metadata.json", metadata))
def _make_content_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
from haiku.rag.utils import escape_sql_string
safe_id = escape_sql_string(did)
rows = list(
client.store.documents_table.search()
.select(["content"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
return rows[0]["content"] if rows else ""
return _run_async(_fetch())
return read_content
def _make_items_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_items(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
items = (
await client.document_item_repository.get_items_in_range(
did, 0, 999999
)
)
lines = []
for item in items:
lines.append(
json.dumps(
{
"position": item.position,
"self_ref": item.self_ref,
"label": item.label,
"text": item.text,
"page_numbers": item.page_numbers,
},
ensure_ascii=False,
)
)
return "\n".join(lines)
return _run_async(_fetch())
return read_items
files.append(
CallbackFile(
f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id),
write=lambda _p, _c: None,
)
)
files.append(
CallbackFile(
f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id),
write=lambda _p, _c: None,
)
)
return OSAccess(files)
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty interpreter."""
external_fns = self._build_external_functions()
vfs = await self._build_vfs()
input_names: list[str] = []
inputs: dict[str, Any] | None = None
@ -166,6 +258,7 @@ class Sandbox:
external_functions=external_fns,
limits=limits,
print_callback=print_callback,
os=vfs,
)
except pydantic_monty.MontyRuntimeError as e:
stdout = "".join(stdout_lines)

218
tests/agents/analysis/test_sandbox.py
View file

@ -74,12 +74,12 @@ class TestSandboxErrors:
assert result.stderr != ""
class TestSandboxHaikuRAG:
"""Test haiku.rag functions in sandbox."""
class TestSandboxListDocuments:
"""Test list_documents function in sandbox."""
@pytest.mark.asyncio
async def test_list_documents_empty(self, sandbox):
"""Test list_documents returns empty list for empty database."""
"""list_documents returns empty list for empty database."""
result = await sandbox.execute(
"docs = await list_documents()\nprint(type(docs).__name__, len(docs))"
)
@ -89,7 +89,7 @@ class TestSandboxHaikuRAG:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_with_data(self, temp_db_path):
"""Test list_documents returns documents when populated."""
"""list_documents returns documents when populated."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
@ -109,6 +109,10 @@ class TestSandboxHaikuRAG:
assert "1" in result.stdout
assert "Test Document" in result.stdout
class TestSandboxSearch:
"""Test search function in sandbox."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_with_data(self, temp_db_path):
@ -158,46 +162,6 @@ class TestSandboxHaikuRAG:
assert "True\nTrue" in result.stdout
assert "list\nlist" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_get_document(self, temp_db_path):
"""Test get_document function."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
f"content = await get_document('{doc.id}')\n"
"print('foxes' in content.lower() if content else 'None')"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_get_document_not_found(self, sandbox):
"""Test get_document returns None for missing document."""
result = await sandbox.execute(
"content = await get_document('nonexistent-id')\nprint(content is None)"
)
assert result.success
assert "True" in result.stdout
class TestSandboxSearchExpandsContext:
"""Test that search() returns expanded results."""
@pytest.mark.asyncio
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_search_returns_expanded_content(self, temp_db_path):
@ -303,13 +267,124 @@ class TestSandboxOutputTruncation:
assert len(result.stdout) < 100
class TestSandboxContextFilter:
"""Test context filter is applied."""
class TestSandboxVFS:
"""Test virtual filesystem for document access."""
@pytest.mark.asyncio
async def test_empty_database_has_no_documents(self, sandbox):
"""Empty database has no document directories."""
result = await sandbox.execute(
"from pathlib import Path\nprint(Path('/documents').exists())"
)
assert result.success
# /documents dir may or may not exist when empty, both are valid
# The key is it doesn't error
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_filter_applied_to_list_documents(self, temp_db_path):
"""Test that context filter is passed to list_documents."""
async def test_iterdir_discovers_documents(self, temp_db_path):
"""Path('/documents').iterdir() lists document directories."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="Test content",
uri="test://doc1",
title="Test Document",
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"dirs = list(Path('/documents').iterdir())\n"
"print(len(dirs))\n"
"print(dirs[0].is_dir())"
)
assert result.success
assert "1" in result.stdout
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_metadata_json(self, temp_db_path):
"""metadata.json contains document title and uri."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Test content",
uri="test://doc1",
title="Test Document",
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
f"meta = json.loads(Path('/documents/{doc.id}/metadata.json').read_text())\n"
"print(meta['title'])\n"
"print(meta['uri'])"
)
assert result.success
assert "Test Document" in result.stdout
assert "test://doc1" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_content_txt(self, temp_db_path):
"""content.txt returns full document text (lazy loaded)."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
f"content = Path('/documents/{doc.id}/content.txt').read_text()\n"
"print('foxes' in content.lower())"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_items_jsonl(self, temp_db_path):
"""items.jsonl returns document items as JSONL (lazy loaded)."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_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)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
f"text = Path('/documents/{doc.id}/items.jsonl').read_text()\n"
"lines = text.strip().split('\\n')\n"
"print(len(lines) > 0)\n"
"item = json.loads(lines[0])\n"
"print('position' in item)\n"
"print('self_ref' in item)\n"
"print('label' in item)\n"
"print('text' in item)\n"
"print('page_numbers' in item)"
)
assert result.success
assert result.stdout.count("True") == 6
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_context_filter_limits_vfs(self, temp_db_path):
"""Context filter restricts which documents appear in VFS."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
@ -326,10 +401,12 @@ class TestSandboxContextFilter:
context = AnalysisContext(filter="uri LIKE 'public://%'")
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
"docs = await list_documents()\n"
"print(len(docs))\n"
"if docs:\n"
" print(docs[0]['title'])"
"from pathlib import Path\n"
"import json\n"
"dirs = list(Path('/documents').iterdir())\n"
"print(len(dirs))\n"
"meta = json.loads((dirs[0] / 'metadata.json').read_text())\n"
"print(meta['title'])"
)
assert result.success
assert "1" in result.stdout
@ -368,43 +445,6 @@ 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 = AnalysisContext()
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."""

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_list_documents_with_data.yaml → tests/cassettes/test_sandbox/TestSandboxListDocuments.test_list_documents_with_data.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_search_returns_doc_item_refs_and_labels.yaml → tests/cassettes/test_sandbox/TestSandboxSearch.test_search_returns_doc_item_refs_and_labels.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_search_with_data.yaml → tests/cassettes/test_sandbox/TestSandboxSearch.test_search_returns_expanded_content.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxSearchExpandsContext.test_search_returns_expanded_content.yaml → tests/cassettes/test_sandbox/TestSandboxSearch.test_search_with_data.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_get_document.yaml → tests/cassettes/test_sandbox/TestSandboxVFS.test_content_txt.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxContextFilter.test_filter_applied_to_list_documents.yaml → tests/cassettes/test_sandbox/TestSandboxVFS.test_context_filter_limits_vfs.yaml
View file

82
tests/cassettes/test_sandbox/TestSandboxVFS.test_items_jsonl.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestSandboxVFS.test_iterdir_discovers_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestSandboxVFS.test_metadata_json.yaml Normal file
View file

File diff suppressed because one or more lines are too long