multimodal sandbox: show_image + picture_refs, drop llm

This commit is contained in:
Yiorgis Gozadinos 2026-05-15 16:13:04 +03:00
parent 2831030f9b
commit 387c4f12a9
No known key found for this signature in database
12 changed files with 308 additions and 1224 deletions

View file

@ -5,6 +5,11 @@
- **`heading_level` and `tree_depth` on `DocumentItem`.** `extract_items` now captures docling's `SectionHeaderItem.level` (H1H6 for headers, `0` elsewhere) and the traversal depth from `iterate_items()` for every item, persisting both in the `document_items` table. Foundations for tree-based document navigation in the analysis sandbox. The 0.48.0 migration adds the columns to existing DBs and backfills them from each doc's docling blob.
- **`toc.json` in the analysis sandbox VFS.** Each document mounted under `/documents/{id}/` now exposes a `toc.json` view alongside `metadata.json`, `content.txt`, `items.jsonl`. Nodes carry `{self_ref, level, title, position, page_numbers, item_range, children}`; `item_range = [start, end_exclusive]` over the same `position` ints used in `items.jsonl`, so the agent can slice items by range to read a section. HTML/markdown ingests produce a real nested tree; PDF ingests produce a flat sibling list because docling collapses heading levels on PDFs. `tree: []` when the doc has no section headers. `items.jsonl` now surfaces `heading_level` and `tree_depth` on every row.
- **Multimodal analysis sandbox.** New `await show_image(document_id, self_ref)` external function inside the analysis sandbox. The picture's bytes are PIL-verified and attached to the `execute_code` tool's response as a pydantic-ai `BinaryContent` part, so a vision-capable driving model sees the actual figure alongside the printed stdout — same mechanism the QA agent's search tool uses. `search()` results now include a `picture_refs` key (subset of `doc_item_refs` labeled `picture`) so the agent can spot which results are figures with one lookup.
### Changed
- **`llm()` removed from the analysis sandbox.** The function was a thin wrapper that spun up an ad-hoc pydantic-ai `Agent` per call. The driving agent already is the LLM — there's no need for a sandbox-internal one. Sandbox external functions are now `search`, `list_documents`, and `show_image`.
### Changed

View file

@ -1,4 +1,5 @@
from pydantic_ai import Agent, RunContext
from pydantic_ai.messages import ToolReturn
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
@ -32,19 +33,26 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisR
)
@agent.tool
async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution:
async def execute_code(
ctx: RunContext[AnalysisDeps], code: str
) -> CodeExecution | ToolReturn:
"""Execute Python code in a sandboxed interpreter.
The code has access to search() and llm() functions, and a
virtual filesystem at /documents/ with document content and structure.
The code has access to search(), list_documents(), and show_image()
external functions, and a virtual filesystem at /documents/ with
document content and structure. Use print() to output results.
Use print() to output results.
When the code calls ``show_image(document_id, self_ref)``, the queued
picture bytes are attached to the tool response as ``BinaryContent``
so a vision-capable driving model can actually see the image.
Args:
code: Python code to execute.
Returns:
Structured result with success status, stdout, and stderr.
Structured result with success status, stdout, and stderr. Wrapped
in a ``ToolReturn`` carrying ``BinaryContent`` parts whenever the
code called ``show_image``.
"""
result = await ctx.deps.sandbox.execute(code)
@ -55,6 +63,10 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisR
success=result.success,
)
if result.binary_attachments:
return ToolReturn(
return_value=execution, content=list(result.binary_attachments)
)
return execution
return agent

View file

@ -12,16 +12,28 @@ 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, doc_item_refs, labels
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs.
`picture_refs` is the subset of `doc_item_refs` whose label is `picture` use it to spot results that contain figures you can surface to the model via `show_image`.
### await list_documents() -> list[dict]
List all documents in the knowledge base.
Returns list of dicts with keys: id, title, uri, created_at
### 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.
### await show_image(document_id, self_ref) -> None
Surface a document picture to the driving LLM as a vision input. The picture's
bytes are attached to this `execute_code` tool's response as a `BinaryContent`
part, so a vision-capable model sees the actual image alongside the printed
output. Missing refs and unverifiable payloads are silent no-ops. Only useful
when the configured analysis model is vision-capable; otherwise the model
receives the bytes but ignores them.
```python
results = await search("revenue chart", limit=5)
for r in results:
for ref in r["picture_refs"]:
await show_image(r["document_id"], ref)
print(f"showed {r['document_id']}:{ref}")
```
## Document Filesystem
@ -155,7 +167,7 @@ Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available)
3b. **Use toc.json for Section Navigation**: When a question is scoped to a section, open `toc.json`, find the matching node, and slice `items.jsonl` by its `item_range` instead of streaming `content.txt`. For PDFs where the tree is flat, the sibling list is still useful as a TOC.
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.
6. **Show pictures explicitly**: When a search result has `picture_refs` and the question is about figures/charts/diagrams, call `show_image(doc_id, ref)` so the driving model can see the picture. Don't dump bytes into stdout.
## Output Format

View file

@ -3,11 +3,12 @@ import atexit
import concurrent.futures
import json
from collections.abc import Callable
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from pydantic_ai import BinaryContent
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
from haiku.rag.agents.analysis.dependencies import AnalysisContext
@ -26,6 +27,7 @@ class SandboxResult:
stdout: str
stderr: str
success: bool
binary_attachments: list[BinaryContent] = field(default_factory=list)
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
@ -95,8 +97,8 @@ class Sandbox:
"""Execute code in a sandboxed Python interpreter.
Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
External functions (search, llm) are called by Monty code using ``await``
and resolved asynchronously on the host.
External functions (search, list_documents, show_image) are called by
Monty code using ``await`` and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
The interpreter uses a REPL session variables persist across
@ -113,6 +115,7 @@ class Sandbox:
_search_results: "list[SearchResult]"
_items_cache: dict[str, str] | None
_toc_cache: dict[str, str] | None
_pending_binary: list[BinaryContent]
_repl: MontyRepl | None
_vfs: OSAccess | None
@ -128,6 +131,7 @@ class Sandbox:
self._search_results = []
self._items_cache = None
self._toc_cache = None
self._pending_binary = []
self._repl = None
self._vfs = None
@ -144,21 +148,29 @@ class Sandbox:
results = await rag.search(query, limit=limit, filter=context.filter)
expanded = await rag.expand_context(results)
self._search_results.extend(expanded)
return [
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
"doc_item_refs": r.doc_item_refs,
"labels": r.labels,
}
for r in expanded
]
out: list[dict[str, Any]] = []
for r in expanded:
picture_refs = [
ref
for ref, lbl in zip(r.doc_item_refs, r.labels, strict=False)
if lbl == "picture"
]
out.append(
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
"doc_item_refs": r.doc_item_refs,
"labels": r.labels,
"picture_refs": picture_refs,
}
)
return out
async def list_documents() -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG
@ -175,20 +187,41 @@ class Sandbox:
for d in docs
]
async def llm(prompt: str) -> str:
from pydantic_ai import Agent
sandbox = self
from haiku.rag.utils import get_model
async def show_image(document_id: str, self_ref: str) -> None:
"""Surface a document picture to the driving LLM as a vision input.
model = get_model(config.analysis.model, config)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
Looks up the picture's bytes by (document_id, self_ref), verifies the
payload via PIL, and queues a ``BinaryContent`` part on the next
``execute_code`` tool return. The driving model sees the picture as
content alongside the textual stdout. Missing refs and unverifiable
payloads are silent no-ops.
"""
from io import BytesIO
from PIL import Image
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
data = await rag.document_item_repository.get_picture_bytes(
document_id, self_ref
)
if not data:
return
try:
Image.open(BytesIO(data)).verify()
except Exception:
return
sandbox._pending_binary.append(
BinaryContent(data=data, media_type="image/png", identifier=self_ref)
)
return {
"search": search,
"list_documents": list_documents,
"llm": llm,
"show_image": show_image,
}
async def _build_vfs(self) -> OSAccess:
@ -376,9 +409,12 @@ class Sandbox:
"""Execute Python code in the Monty REPL.
Variables persist across calls within the same Sandbox instance.
Binary attachments queued by ``show_image`` during this call are
drained into the returned ``SandboxResult.binary_attachments``.
"""
repl, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions()
self._pending_binary = []
stdout_lines: list[str] = []
@ -401,7 +437,12 @@ class Sandbox:
stdout = "".join(stdout_lines)
if len(stdout) > max_chars:
stdout = stdout[:max_chars] + "\n... (output truncated)"
return SandboxResult(stdout=stdout, stderr=str(e), success=False)
return SandboxResult(
stdout=stdout,
stderr=str(e),
success=False,
binary_attachments=self._pending_binary,
)
stdout = "".join(stdout_lines)
if output is not None:
@ -414,4 +455,9 @@ class Sandbox:
stdout_with_output[:max_chars] + "\n... (output truncated)"
)
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)
return SandboxResult(
stdout=stdout_with_output,
stderr="",
success=True,
binary_attachments=self._pending_binary,
)

View file

@ -26,7 +26,7 @@ Retrieve a document by ID, title, or URI. Partial matches work.
{% if "execute_code" in tool_names %}
### execute_code
Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, `await llm()`, and a virtual filesystem at `/documents/` with document content and structure.
Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, `await show_image(document_id, self_ref)`, and a virtual filesystem at `/documents/` with document content and structure.
{% endif %}
{% if "cite" in tool_names %}

View file

@ -242,9 +242,10 @@ def create_skill_tools(
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter.
The code has access to search(), list_documents(), llm() functions
and a virtual filesystem at /documents/ with document content and
structure (metadata.json, content.txt, items.jsonl per document).
The code has access to search(), list_documents(), show_image()
functions and a virtual filesystem at /documents/ with document
content and structure (metadata.json, content.txt, items.jsonl,
toc.json per document).
Use print() to output results. Variables persist between calls
within the same skill invocation.

View file

@ -18,9 +18,9 @@ You solve complex analytical questions by writing and executing Python code agai
Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results.
Inside the code, these functions are available (use `await`):
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`)
- `await list_documents()` → list of dicts with keys: id, title, uri, created_at
- `await llm(prompt)` → string response from an LLM (for classification, summarization, extraction)
- `await show_image(document_id, self_ref)` → attaches a document picture's bytes as a `BinaryContent` part on this tool's response so a vision-capable model sees it. Silent no-op for missing or unverifiable refs.
Available modules: `json`, `re`, `math`, `pathlib`
Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
@ -102,6 +102,6 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
- Variables persist between `execute_code` calls — you can search in one call and process results in the next
- Use `print()` to output results — the output is your only feedback
- Always execute code to answer questions — don't just describe what code would do
- Use `await` for all async functions inside execute_code (search, list_documents, llm)
- Use `await` for all async functions inside execute_code (search, list_documents, show_image)
- Use `Path.read_text()` to read files — do NOT use `open()`, `with` statements, or `collections` module
- Do NOT include chunk IDs or UUIDs in your answer text — use the `cite` tool separately

View file

@ -160,59 +160,6 @@ class TestClientAnalysisIntegration:
assert "Animal Facts" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_semantic_analysis_with_llm(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can use llm() for semantic analysis combined with computation.
Agent program:
docs = list_documents(limit=100)
print(len(docs))
print([d['title'] for d in docs[:20]])
sentiments = {}
for title in ['Q1 Update', 'Q2 Update', 'Q3 Update']:
content = get_document(title)
if content:
result = llm(f"Classify sentiment as positive/negative/mixed: {content}")
sentiments[title] = result
print(sentiments)
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"The new product launch exceeded expectations. Sales grew 40% "
"and customer feedback has been overwhelmingly positive. "
"Team morale is at an all-time high.",
title="Q1 Update",
)
await client.create_document(
"We faced significant challenges this quarter. Supply chain issues "
"caused delays, and we missed our revenue target by 15%. "
"Several key employees left the company.",
title="Q2 Update",
)
await client.create_document(
"Mixed results this quarter. While product quality improved, "
"marketing campaigns underperformed. Revenue was flat compared "
"to last year but customer retention increased.",
title="Q3 Update",
)
result = await client.analyze(
"Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?"
)
# Should identify: Q1=positive, Q2=negative, Q3=mixed
assert "positive" in result.answer.lower()
assert "negative" in result.answer.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):

View file

@ -446,22 +446,3 @@ class TestSandboxPreloadedDocuments:
assert "2" in result.stdout
assert "Doc A" in result.stdout
assert "Doc B" in result.stdout
class TestSandboxLLM:
"""Test llm() external function."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_llm_function(self, allow_model_requests, temp_db_path):
"""Test llm() calls the model and returns a string."""
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"answer = await llm('What is 2 + 2? Reply with just the number.')\n"
"print(answer)"
)
assert result.success
assert "4" in result.stdout

View file

@ -0,0 +1,186 @@
"""Tests for the multimodal sandbox surface: show_image, picture_refs in
search results, binary_attachments on SandboxResult, and the absence of
the old llm() external function.
"""
import base64
from io import BytesIO
import pytest
from PIL import Image
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
def _png_bytes(color: str = "red", size: tuple[int, int] = (8, 8)) -> bytes:
"""Generate a real PNG so PIL.Image.verify() accepts it."""
img = Image.new("RGB", size, color)
buf = BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
async def _seed_doc_with_picture(client, *, png: bytes) -> tuple[str, str]:
"""Create a Document row, replace its items with one picture row carrying
the given bytes. Returns (doc_id, self_ref)."""
doc = await client.create_document(content="x", uri="test://pic", title="Pic")
await client.document_item_repository.delete_by_document_id(doc.id)
self_ref = "#/pictures/0"
items = [
DocumentItem(
document_id=doc.id,
position=0,
self_ref=self_ref,
label="picture",
text="",
page_numbers=[1],
picture_data=png,
)
]
await client.document_item_repository.create_items(doc.id, items)
return doc.id, self_ref
@pytest.mark.asyncio
class TestShowImage:
"""show_image() appends a BinaryContent attachment when bytes verify."""
async def test_appends_binary_attachment(self, temp_db_path):
png = _png_bytes("red")
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id, ref = await _seed_doc_with_picture(client, png=png)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
result = await sandbox.execute(
f"await show_image({doc_id!r}, {ref!r})\nprint('done')"
)
assert result.success, result.stderr
assert len(result.binary_attachments) == 1
att = result.binary_attachments[0]
assert att.media_type == "image/png"
assert att.identifier == ref
assert att.data == png
async def test_missing_picture_is_silent_noop(self, temp_db_path):
png = _png_bytes("red")
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id, _ = await _seed_doc_with_picture(client, png=png)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
result = await sandbox.execute(
f"await show_image({doc_id!r}, '#/pictures/999')\nprint('ok')"
)
assert result.success, result.stderr
assert result.binary_attachments == []
async def test_invalid_bytes_rejected(self, temp_db_path):
# Garbage bytes — PIL.verify() should refuse, no attachment emitted.
garbage = b"this is not a PNG"
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id, ref = await _seed_doc_with_picture(client, png=garbage)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
result = await sandbox.execute(
f"await show_image({doc_id!r}, {ref!r})\nprint('checked')"
)
assert result.success, result.stderr
assert result.binary_attachments == []
async def test_attachments_reset_across_executes(self, temp_db_path):
png = _png_bytes("red")
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id, ref = await _seed_doc_with_picture(client, png=png)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
first = await sandbox.execute(
f"await show_image({doc_id!r}, {ref!r})\nprint('first')"
)
second = await sandbox.execute("print('second')")
assert first.success and len(first.binary_attachments) == 1
assert second.success and second.binary_attachments == []
@pytest.mark.asyncio
class TestSearchPictureRefs:
"""search() result dicts carry a `picture_refs` list (subset of
doc_item_refs labeled 'picture'). No `image_data` base64 in the dict."""
async def test_picture_refs_extracted_from_labels(self, temp_db_path, monkeypatch):
# Build a fake SearchResult with mixed labels so we don't need an embedder.
synthetic = [
SearchResult(
chunk_id="c1",
content="hit",
document_id="d1",
document_uri="test://d1",
document_title="D1",
score=1.0,
page_numbers=[1],
headings=None,
doc_item_refs=["#/texts/0", "#/pictures/0", "#/pictures/1"],
labels=["text", "picture", "picture"],
),
SearchResult(
chunk_id="c2",
content="text only",
document_id="d1",
document_uri="test://d1",
document_title="D1",
score=0.5,
page_numbers=[2],
headings=None,
doc_item_refs=["#/texts/5"],
labels=["text"],
),
]
async def fake_search(self, *args, **kwargs):
return synthetic
async def fake_expand_context(self, results):
return results
# Patch HaikuRAG.search and expand_context so the sandbox closure runs
# without an embedder. The sandbox opens its own HaikuRAG instance, so
# we patch on the class.
monkeypatch.setattr(HaikuRAG, "search", fake_search)
monkeypatch.setattr(HaikuRAG, "expand_context", fake_expand_context)
async with HaikuRAG(temp_db_path, create=True):
pass # ensure the DB exists so the sandbox can open it read-only
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
external = sandbox._build_external_functions()
results = await external["search"]("anything")
assert len(results) == 2
assert results[0]["picture_refs"] == ["#/pictures/0", "#/pictures/1"]
assert results[1]["picture_refs"] == []
# No raw base64 garbage in the dict.
assert "image_data" not in results[0]
@pytest.mark.asyncio
class TestExternalFunctionsShape:
"""llm() is gone; show_image() is present."""
async def test_llm_gone_show_image_present(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True):
pass
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
external = sandbox._build_external_functions()
assert "llm" not in external
assert "show_image" in external
assert "search" in external
assert "list_documents" in external
# Silence unused-import flake — base64 is reserved for follow-up tests that
# decode attachment.data and compare. Kept eagerly imported for parity with
# the QA binary-content tests.
_ = base64

View file

@ -1,51 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '143'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: What is 2 + 2? Reply with just the number.
role: user
model: gpt-oss
reasoning_effort: low
stream: false
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '311'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: '4'
reasoning: Just reply 4.
role: assistant
created: 1771924616
id: chatcmpl-525
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 16
prompt_tokens: 81
total_tokens: 97
status:
code: 200
message: OK
version: 1