multimodal sandbox: show_image + picture_refs, drop llm
This commit is contained in:
parent
2831030f9b
commit
387c4f12a9
12 changed files with 308 additions and 1224 deletions
|
|
@ -5,6 +5,11 @@
|
|||
|
||||
- **`heading_level` and `tree_depth` on `DocumentItem`.** `extract_items` now captures docling's `SectionHeaderItem.level` (H1–H6 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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 %}
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
186
tests/agents/analysis/test_sandbox_multimodal.py
Normal file
186
tests/agents/analysis/test_sandbox_multimodal.py
Normal 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
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -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
|
||||
Loading…
Reference in a new issue