route figures through cite, drop show_image
This commit is contained in:
Yiorgis Gozadinos
parent
387c4f12a9
commit
f13461cdca
14 changed files with 174 additions and 239 deletions
|
|
@ -5,11 +5,12 @@
|
|||
|
||||
- **`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.
|
||||
- **`picture_refs` in sandbox `search()` results.** Each search dict in the analysis sandbox now exposes a `picture_refs` key (subset of `doc_item_refs` whose label is `picture`) so the agent can spot picture chunks without zipping refs + labels manually.
|
||||
- **Citations carry `picture_refs` for inline picture rendering.** `Citation` gains a `picture_refs: list[str]` field; `resolve_citations` derives it from the originating `SearchResult.doc_item_refs` + `labels`. The chat TUI's `CitationWidget` mounts a `textual_image.widget.Image` per picture self_ref inside the existing collapsible — picture citations render the actual figure alongside their text content. The skill's search tool already auto-attaches `BinaryContent` for picture chunks under vision-capable models, so the driving model sees the figure during reasoning; the citation pipeline now also surfaces it in the user's UI. `visualize_chunk` continues to work uniformly because every cited chunk has a `chunk_id`.
|
||||
|
||||
### 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`.
|
||||
- **`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` and `list_documents`.
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
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
|
||||
|
|
@ -33,40 +32,25 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisR
|
|||
)
|
||||
|
||||
@agent.tool
|
||||
async def execute_code(
|
||||
ctx: RunContext[AnalysisDeps], code: str
|
||||
) -> CodeExecution | ToolReturn:
|
||||
async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution:
|
||||
"""Execute Python code in a sandboxed interpreter.
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
The code has access to search() and list_documents() external
|
||||
functions, plus a virtual filesystem at /documents/ with document
|
||||
content and structure. Use print() to output results.
|
||||
|
||||
Args:
|
||||
code: Python code to execute.
|
||||
|
||||
Returns:
|
||||
Structured result with success status, stdout, and stderr. Wrapped
|
||||
in a ``ToolReturn`` carrying ``BinaryContent`` parts whenever the
|
||||
code called ``show_image``.
|
||||
Structured result with success status, stdout, and stderr.
|
||||
"""
|
||||
result = await ctx.deps.sandbox.execute(code)
|
||||
|
||||
execution = CodeExecution(
|
||||
return CodeExecution(
|
||||
code=code,
|
||||
stdout=result.stdout,
|
||||
stderr=result.stderr,
|
||||
success=result.success,
|
||||
)
|
||||
|
||||
if result.binary_attachments:
|
||||
return ToolReturn(
|
||||
return_value=execution, content=list(result.binary_attachments)
|
||||
)
|
||||
return execution
|
||||
|
||||
return agent
|
||||
|
|
|
|||
|
|
@ -13,28 +13,12 @@ Inside execute_code, these functions are ALREADY available in the namespace. Do
|
|||
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, 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`.
|
||||
`picture_refs` is the subset of `doc_item_refs` whose label is `picture` — use it to spot results that contain figures.
|
||||
|
||||
### await list_documents() -> list[dict]
|
||||
List all documents in the knowledge base.
|
||||
Returns list of dicts with keys: id, title, uri, created_at
|
||||
|
||||
### 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
|
||||
|
||||
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.
|
||||
|
|
@ -167,7 +151,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. **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.
|
||||
6. **Cite picture chunks for figure-driven questions**: When a question is about a figure or diagram, find the picture chunk (search results with non-empty `picture_refs`) and cite its chunk_id. The driving model already sees figures from search hits; the citation makes the picture visible in the user's UI as well.
|
||||
|
||||
## Output Format
|
||||
|
||||
|
|
|
|||
|
|
@ -3,12 +3,11 @@ import atexit
|
|||
import concurrent.futures
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
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
|
||||
|
|
@ -27,7 +26,6 @@ class SandboxResult:
|
|||
stdout: str
|
||||
stderr: str
|
||||
success: bool
|
||||
binary_attachments: list[BinaryContent] = field(default_factory=list)
|
||||
|
||||
|
||||
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
|
||||
|
|
@ -63,8 +61,6 @@ def _build_toc(items: list["DocumentItem"]) -> list[dict[str, Any]]:
|
|||
|
||||
total = max((i.position for i in items), default=-1) + 1
|
||||
|
||||
# Pre-compute each header's end position: the next header in document order
|
||||
# whose heading_level <= this header's level.
|
||||
ends: list[int] = []
|
||||
for idx, h in enumerate(headers):
|
||||
end = total
|
||||
|
|
@ -97,9 +93,9 @@ 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, 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}/``.
|
||||
External functions (search, list_documents) 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
|
||||
``execute()`` calls within the same Sandbox instance.
|
||||
|
|
@ -115,7 +111,6 @@ 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
|
||||
|
||||
|
|
@ -131,7 +126,6 @@ class Sandbox:
|
|||
self._search_results = []
|
||||
self._items_cache = None
|
||||
self._toc_cache = None
|
||||
self._pending_binary = []
|
||||
self._repl = None
|
||||
self._vfs = None
|
||||
|
||||
|
|
@ -151,9 +145,7 @@ class Sandbox:
|
|||
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"
|
||||
ref for ref in r.doc_item_refs if ref.startswith("#/pictures/")
|
||||
]
|
||||
out.append(
|
||||
{
|
||||
|
|
@ -187,41 +179,9 @@ class Sandbox:
|
|||
for d in docs
|
||||
]
|
||||
|
||||
sandbox = self
|
||||
|
||||
async def show_image(document_id: str, self_ref: str) -> None:
|
||||
"""Surface a document picture to the driving LLM as a vision input.
|
||||
|
||||
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,
|
||||
"show_image": show_image,
|
||||
}
|
||||
|
||||
async def _build_vfs(self) -> OSAccess:
|
||||
|
|
@ -230,7 +190,8 @@ class Sandbox:
|
|||
Mounts per-document directories with:
|
||||
- metadata.json: MemoryFile (eager, small)
|
||||
- content.txt: CallbackFile (lazy, can be large)
|
||||
- items.jsonl: CallbackFile (lazy, can be large)
|
||||
- items.jsonl: CallbackFile (lazy, bulk-cached)
|
||||
- toc.json: CallbackFile (lazy, bulk-cached)
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
|
@ -409,12 +370,9 @@ 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] = []
|
||||
|
||||
|
|
@ -437,12 +395,7 @@ 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,
|
||||
binary_attachments=self._pending_binary,
|
||||
)
|
||||
return SandboxResult(stdout=stdout, stderr=str(e), success=False)
|
||||
|
||||
stdout = "".join(stdout_lines)
|
||||
if output is not None:
|
||||
|
|
@ -455,9 +408,4 @@ class Sandbox:
|
|||
stdout_with_output[:max_chars] + "\n... (output truncated)"
|
||||
)
|
||||
|
||||
return SandboxResult(
|
||||
stdout=stdout_with_output,
|
||||
stderr="",
|
||||
success=True,
|
||||
binary_attachments=self._pending_binary,
|
||||
)
|
||||
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,11 @@ class Citation(BaseModel):
|
|||
|
||||
Used by research graph and chat applications. The optional index field
|
||||
supports UI display ordering in chat contexts.
|
||||
|
||||
``picture_refs`` lists the ``self_ref`` values of picture items in the
|
||||
cited chunk. Empty for text-only citations. UIs can fetch the picture
|
||||
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
|
||||
and render them alongside the text content.
|
||||
"""
|
||||
|
||||
index: int | None = None
|
||||
|
|
@ -33,6 +38,7 @@ class Citation(BaseModel):
|
|||
page_numbers: list[int] = Field(default_factory=list)
|
||||
headings: list[str] | None = None
|
||||
content: str
|
||||
picture_refs: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RawSearchAnswer(BaseModel):
|
||||
|
|
@ -98,6 +104,7 @@ def resolve_citations(
|
|||
r = by_id.get(chunk_id)
|
||||
if not r:
|
||||
continue
|
||||
picture_refs = [ref for ref in r.doc_item_refs if ref.startswith("#/pictures/")]
|
||||
citations.append(
|
||||
Citation(
|
||||
document_id=r.document_id or "",
|
||||
|
|
@ -107,6 +114,7 @@ def resolve_citations(
|
|||
page_numbers=r.page_numbers,
|
||||
headings=r.headings,
|
||||
content=r.content,
|
||||
picture_refs=picture_refs,
|
||||
)
|
||||
)
|
||||
return citations
|
||||
|
|
|
|||
|
|
@ -339,8 +339,26 @@ class ChatApp(App):
|
|||
for cid in cited_ids:
|
||||
if cid in citation_index:
|
||||
citations.append(citation_index[cid])
|
||||
if citations:
|
||||
await chat_history.add_citations(citations)
|
||||
if not citations:
|
||||
return
|
||||
|
||||
picture_bytes: dict[str, list[bytes]] = {}
|
||||
if self.client is not None:
|
||||
for citation in citations:
|
||||
refs = list(citation.picture_refs or [])
|
||||
if not refs:
|
||||
continue
|
||||
blobs: list[bytes] = []
|
||||
for ref in refs:
|
||||
data = await self.client.document_item_repository.get_picture_bytes(
|
||||
citation.document_id, ref
|
||||
)
|
||||
if data:
|
||||
blobs.append(data)
|
||||
if blobs:
|
||||
picture_bytes[citation.chunk_id] = blobs
|
||||
|
||||
await chat_history.add_citations(citations, picture_bytes=picture_bytes)
|
||||
|
||||
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
|
||||
if analysis_state:
|
||||
|
|
|
|||
|
|
@ -1,10 +1,13 @@
|
|||
from io import BytesIO
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from PIL import Image as PILImage
|
||||
from textual.containers import Horizontal, VerticalScroll
|
||||
from textual.css.query import NoMatches
|
||||
from textual.message import Message
|
||||
from textual.widgets import Collapsible, LoadingIndicator, Markdown, Static
|
||||
from textual.widgets.markdown import MarkdownStream
|
||||
from textual_image.widget import Image as TextualImage
|
||||
|
||||
from haiku.rag.agents.research.models import Citation
|
||||
|
||||
|
|
@ -119,7 +122,12 @@ class CitationWidget(Collapsible):
|
|||
super().__init__()
|
||||
self.widget = widget
|
||||
|
||||
def __init__(self, citation: Citation, **kwargs) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
citation: Citation,
|
||||
picture_bytes: list[bytes] | None = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
title = f"[{citation.index}] {citation.document_title or citation.document_uri}"
|
||||
if citation.page_numbers:
|
||||
pages = ", ".join(map(str, citation.page_numbers[:3]))
|
||||
|
|
@ -131,7 +139,13 @@ class CitationWidget(Collapsible):
|
|||
if len(content) > 500:
|
||||
content = content[:500] + "..."
|
||||
|
||||
children: list[Markdown | Static] = [Markdown(content)]
|
||||
children: list[Any] = [Markdown(content)]
|
||||
for blob in picture_bytes or []:
|
||||
try:
|
||||
pil = PILImage.open(BytesIO(blob))
|
||||
except Exception:
|
||||
continue
|
||||
children.append(TextualImage(pil, classes="citation-image"))
|
||||
if citation.headings:
|
||||
headings = " > ".join(citation.headings[:3])
|
||||
children.append(Static(f"Section: {headings}", classes="citation-metadata"))
|
||||
|
|
@ -333,6 +347,14 @@ class ChatHistory(VerticalScroll):
|
|||
text-style: italic;
|
||||
}
|
||||
|
||||
CitationWidget .citation-image {
|
||||
width: auto;
|
||||
height: auto;
|
||||
max-width: 100%;
|
||||
max-height: 30;
|
||||
margin: 1 0;
|
||||
}
|
||||
|
||||
/* Program */
|
||||
ProgramWidget {
|
||||
margin: 0 0 0 2;
|
||||
|
|
@ -414,13 +436,26 @@ class ChatHistory(VerticalScroll):
|
|||
widget.mark_completed()
|
||||
widget.add_class("complete")
|
||||
|
||||
async def add_citations(self, citations: list[Citation]) -> None:
|
||||
"""Add citations inline after a response."""
|
||||
async def add_citations(
|
||||
self,
|
||||
citations: list[Citation],
|
||||
picture_bytes: dict[str, list[bytes]] | None = None,
|
||||
) -> None:
|
||||
"""Add citations inline after a response.
|
||||
|
||||
``picture_bytes`` maps citation ``chunk_id`` → list of raw PNG bytes,
|
||||
one per entry in the citation's ``picture_refs``. Pre-fetched by the
|
||||
caller (typically the chat app's post-response hook) so widget
|
||||
construction stays synchronous.
|
||||
"""
|
||||
if not citations:
|
||||
return
|
||||
await self.mount(SourcesHeader(len(citations)))
|
||||
picture_bytes = picture_bytes or {}
|
||||
for citation in citations:
|
||||
widget = CitationWidget(citation)
|
||||
widget = CitationWidget(
|
||||
citation, picture_bytes=picture_bytes.get(citation.chunk_id)
|
||||
)
|
||||
await self.mount(widget)
|
||||
self.scroll_end(animate=False)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 show_image(document_id, self_ref)`, 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()`, and a virtual filesystem at `/documents/` with document content and structure.
|
||||
{% endif %}
|
||||
{% if "cite" in tool_names %}
|
||||
|
||||
|
|
|
|||
|
|
@ -242,10 +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(), show_image()
|
||||
functions and a virtual filesystem at /documents/ with document
|
||||
content and structure (metadata.json, content.txt, items.jsonl,
|
||||
toc.json per document).
|
||||
The code has access to search() and list_documents() 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.
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ Execute Python code in a sandboxed interpreter. Variables persist between calls
|
|||
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, picture_refs (subset of doc_item_refs labeled `picture`)
|
||||
- `await list_documents()` → list of dicts with keys: id, title, uri, created_at
|
||||
- `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
|
||||
|
|
@ -32,7 +31,7 @@ Search the knowledge base directly (outside code execution). Use for initial exp
|
|||
List available documents. Use to discover what's in the knowledge base.
|
||||
|
||||
### cite
|
||||
Register chunk IDs as citations. Call after your analysis with chunk_id values from search results that support your answer.
|
||||
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results (from either the `search` tool or `await search(...)` inside `execute_code`) that support each claim. Every answer that uses search results must be backed by `cite`.
|
||||
|
||||
## Document Filesystem (inside execute_code)
|
||||
|
||||
|
|
@ -95,13 +94,16 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
|
|||
1. Use `search` tool first to understand what's in the knowledge base
|
||||
2. Use `execute_code` to write analysis code
|
||||
3. Iterate: run code, examine output, refine approach
|
||||
4. Call `cite` with chunk IDs from search results you referenced
|
||||
4. Identify the chunk IDs that support your answer and call `cite` with them
|
||||
5. Then write a concise answer based strictly on the cited content
|
||||
|
||||
You MUST call `cite` with at least one chunk ID before producing your final answer, **unless** you are refusing for lack of information. Answers without citations are considered ungrounded. In a refusal case do **not** call `cite` — there is nothing to cite.
|
||||
|
||||
## Important
|
||||
|
||||
- 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, show_image)
|
||||
- Use `await` for all async functions inside execute_code (search, list_documents)
|
||||
- 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
|
||||
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
|
||||
|
|
|
|||
|
|
@ -1,108 +1,14 @@
|
|||
"""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.
|
||||
"""Tests for the multimodal sandbox surface: picture_refs on search dicts
|
||||
and the set of external functions exposed to the interpreter.
|
||||
"""
|
||||
|
||||
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
|
||||
|
|
@ -110,8 +16,12 @@ 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.
|
||||
async def test_picture_refs_extracted_from_self_refs(
|
||||
self, temp_db_path, monkeypatch
|
||||
):
|
||||
# ``labels`` is a deduplicated set after expand_context and is not
|
||||
# aligned with ``doc_item_refs``, so picture refs must be recovered
|
||||
# from the docling ``#/pictures/...`` self_ref convention.
|
||||
synthetic = [
|
||||
SearchResult(
|
||||
chunk_id="c1",
|
||||
|
|
@ -123,7 +33,7 @@ class TestSearchPictureRefs:
|
|||
page_numbers=[1],
|
||||
headings=None,
|
||||
doc_item_refs=["#/texts/0", "#/pictures/0", "#/pictures/1"],
|
||||
labels=["text", "picture", "picture"],
|
||||
labels=["picture", "text"],
|
||||
),
|
||||
SearchResult(
|
||||
chunk_id="c2",
|
||||
|
|
@ -145,9 +55,6 @@ class TestSearchPictureRefs:
|
|||
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)
|
||||
|
||||
|
|
@ -161,26 +68,22 @@ class TestSearchPictureRefs:
|
|||
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."""
|
||||
"""Sandbox externals expose exactly ``search`` and ``list_documents``.
|
||||
|
||||
async def test_llm_gone_show_image_present(self, temp_db_path):
|
||||
Pictures reach the driving model through the skill's own ``search``
|
||||
tool (BinaryContent auto-attach for picture chunks) and surface in
|
||||
the UI as citations with ``picture_refs`` populated by
|
||||
``resolve_citations``.
|
||||
"""
|
||||
|
||||
async def test_externals_are_search_and_list_documents_only(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
|
||||
assert set(external) == {"search", "list_documents"}
|
||||
|
|
|
|||
|
|
@ -167,7 +167,7 @@ class TestTocShape:
|
|||
assert toc["tree"] == []
|
||||
|
||||
async def test_skip_header_with_zero_level(self, temp_db_path):
|
||||
"""A section_header with heading_level=0 (legacy pre-0.46.0 row) is skipped."""
|
||||
"""A section_header with ``heading_level == 0`` is skipped."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
doc_id = await _empty_doc(client, uri="test://zero", title="Zero Level")
|
||||
items = [
|
||||
|
|
@ -240,13 +240,11 @@ class TestItemsJsonlSurfacesNewFields:
|
|||
rows = await _read_items_jsonl(sandbox, doc_id)
|
||||
|
||||
assert len(rows) == 3
|
||||
# Field presence + values
|
||||
assert rows[0]["heading_level"] == 1
|
||||
assert rows[0]["tree_depth"] == 2
|
||||
assert rows[1]["heading_level"] == 0
|
||||
assert rows[1]["tree_depth"] == 3
|
||||
assert rows[2]["heading_level"] == 2
|
||||
assert rows[2]["tree_depth"] == 4
|
||||
# Existing fields still present and unchanged
|
||||
for r in rows:
|
||||
assert {"position", "self_ref", "label", "text", "page_numbers"} <= set(r)
|
||||
|
|
|
|||
|
|
@ -175,6 +175,61 @@ class TestExecuteCodeTool:
|
|||
assert "_sandbox" in state.searches
|
||||
assert len(state.searches["_sandbox"]) > 0
|
||||
|
||||
async def test_cite_picture_chunk_records_picture_refs(self, rag_db):
|
||||
"""Citing a picture chunk populates ``Citation.picture_refs`` from the
|
||||
docling ``#/pictures/...`` self_ref prefix.
|
||||
|
||||
``labels`` is a deduplicated set after expand_context and is not
|
||||
aligned with ``doc_item_refs``, so the prefix is the only reliable
|
||||
signal.
|
||||
"""
|
||||
from haiku.rag.skills.analysis import AnalysisState, create_skill
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
|
||||
expanded_picture_hit = SearchResult(
|
||||
chunk_id="picture-chunk-abc",
|
||||
content="picture + surrounding prose",
|
||||
document_id="doc-1",
|
||||
document_uri="test://doc-1",
|
||||
document_title="Doc 1",
|
||||
score=1.0,
|
||||
page_numbers=[1, 2],
|
||||
headings=None,
|
||||
doc_item_refs=[
|
||||
"#/texts/3",
|
||||
"#/pictures/0",
|
||||
"#/texts/4",
|
||||
"#/pictures/1",
|
||||
],
|
||||
labels=["caption", "picture", "text"],
|
||||
)
|
||||
text_result = SearchResult(
|
||||
chunk_id="text-chunk-xyz",
|
||||
content="prose",
|
||||
document_id="doc-1",
|
||||
document_uri="test://doc-1",
|
||||
document_title="Doc 1",
|
||||
score=0.7,
|
||||
page_numbers=[1],
|
||||
headings=None,
|
||||
doc_item_refs=["#/texts/4"],
|
||||
labels=["text"],
|
||||
)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
cite = _get_tool(skill, "cite")
|
||||
state = AnalysisState()
|
||||
state.searches["q1"] = [expanded_picture_hit, text_result]
|
||||
ctx = _make_ctx(state)
|
||||
|
||||
await cite(ctx, chunk_ids=["picture-chunk-abc", "text-chunk-xyz"])
|
||||
|
||||
assert state.citations == ["picture-chunk-abc", "text-chunk-xyz"]
|
||||
picture_citation = state.citation_index["picture-chunk-abc"]
|
||||
assert picture_citation.picture_refs == ["#/pictures/0", "#/pictures/1"]
|
||||
text_citation = state.citation_index["text-chunk-xyz"]
|
||||
assert text_citation.picture_refs == []
|
||||
|
||||
async def test_execute_code_vfs_write_denied(self, rag_db, sandbox_factory):
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
|
|
|
|||
|
|
@ -169,7 +169,6 @@ class TestV0_48_0FreshSchema:
|
|||
assert "tree_depth" in names
|
||||
|
||||
|
||||
# Sanity: ensure module imports without side effects.
|
||||
def test_module_imports():
|
||||
from haiku.rag.store.upgrades import v0_48_0 # noqa: F401
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue