472 lines
18 KiB
Python
472 lines
18 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from collections.abc import AsyncIterator, Callable, Coroutine
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
import pydantic_monty
|
|
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
|
|
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.sandbox.dependencies import AnalysisContext
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import PurePosixPath
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
|
|
@dataclass
|
|
class SandboxResult:
|
|
"""Result of executing code in the sandbox."""
|
|
|
|
stdout: str
|
|
stderr: str
|
|
success: bool
|
|
|
|
|
|
def _build_toc(
|
|
items: list["DocumentItem"],
|
|
chunk_index: dict[str, list[str]],
|
|
) -> list[dict[str, Any]]:
|
|
"""Build a nested section tree from items in position order.
|
|
|
|
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting
|
|
follows the explicit levels: a header pops the stack until the top is at
|
|
a strictly shallower level, then becomes a child of that top (or a root).
|
|
|
|
``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the
|
|
position of the next header whose level is the same or shallower (i.e.
|
|
the next sibling or ancestor that ends this section), or the total item
|
|
count if no such header exists.
|
|
|
|
``chunk_ids`` aggregates the chunks covered by all items in the section's
|
|
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
|
|
ground a section-scoped answer without a corpus-wide ``search()`` call.
|
|
|
|
Items without a section_header label (or with ``heading_level == 0``) are
|
|
skipped. When all section_headers carry the same level the output is a
|
|
flat sibling list (see docling-project/docling#2121 for an upstream case
|
|
where every PDF section_header is emitted at level=1).
|
|
"""
|
|
# Defensive: every consumer is supposed to pass items in position order,
|
|
# but the end_exclusive lookahead below silently miscomputes section
|
|
# boundaries if it's not — better to sort once than trust the caller.
|
|
items = sorted(items, key=lambda i: i.position)
|
|
headers: list[DocumentItem] = [
|
|
i for i in items if i.label == "section_header" and i.heading_level > 0
|
|
]
|
|
if not headers:
|
|
return []
|
|
|
|
total = max((i.position for i in items), default=-1) + 1
|
|
items_by_position: dict[int, DocumentItem] = {i.position: i for i in items}
|
|
|
|
ends: list[int] = []
|
|
for idx, h in enumerate(headers):
|
|
end = total
|
|
for j in range(idx + 1, len(headers)):
|
|
if headers[j].heading_level <= h.heading_level:
|
|
end = headers[j].position
|
|
break
|
|
ends.append(end)
|
|
|
|
roots: list[dict[str, Any]] = []
|
|
stack: list[tuple[int, dict[str, Any]]] = []
|
|
for h, end in zip(headers, ends, strict=True):
|
|
seen: set[str] = set()
|
|
chunk_ids: list[str] = []
|
|
for pos in range(h.position, end):
|
|
item = items_by_position.get(pos)
|
|
if item is None:
|
|
continue
|
|
for cid in chunk_index.get(item.self_ref, []):
|
|
if cid not in seen:
|
|
seen.add(cid)
|
|
chunk_ids.append(cid)
|
|
node: dict[str, Any] = {
|
|
"self_ref": h.self_ref,
|
|
"level": h.heading_level,
|
|
"title": h.text,
|
|
"page_numbers": list(h.page_numbers),
|
|
"item_range": [h.position, end],
|
|
"chunk_ids": chunk_ids,
|
|
"children": [],
|
|
}
|
|
while stack and stack[-1][0] >= h.heading_level:
|
|
stack.pop()
|
|
(stack[-1][1]["children"] if stack else roots).append(node)
|
|
stack.append((h.heading_level, node))
|
|
return roots
|
|
|
|
|
|
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) 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.
|
|
|
|
sandbox = Sandbox(db_path, config, context)
|
|
result = await sandbox.execute("x = await search('query')")
|
|
result = await sandbox.execute("print(x[0]['content'])") # x persists
|
|
|
|
All database access runs on the event loop that drives ``execute()``. Monty's
|
|
file callbacks are synchronous and run on the interpreter's worker thread, so
|
|
they bridge back to that loop via ``run_coroutine_threadsafe``; the loop is
|
|
free during ``feed_run_async`` (the VM runs on the worker thread), so the
|
|
bridge does not deadlock. When a ``rag`` connection is supplied it is used
|
|
for every read, so an analysis run drives a single connection on a single
|
|
loop; otherwise each read opens an ephemeral read-only connection.
|
|
"""
|
|
|
|
_db_path: Path
|
|
_config: AppConfig
|
|
_context: AnalysisContext
|
|
_rag: "HaikuRAG | None"
|
|
_lock: "asyncio.Lock | None"
|
|
_search_results: "list[SearchResult]"
|
|
_doc_items: dict[str, list["DocumentItem"]]
|
|
_doc_chunk_index: dict[str, dict[str, list[str]]]
|
|
_items_jsonl_cache: dict[str, str]
|
|
_toc_json_cache: dict[str, str]
|
|
_repl: MontyRepl | None
|
|
_vfs: OSAccess | None
|
|
_loop: asyncio.AbstractEventLoop | None
|
|
|
|
def __init__(
|
|
self,
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
context: AnalysisContext,
|
|
rag: "HaikuRAG | None" = None,
|
|
lock: "asyncio.Lock | None" = None,
|
|
):
|
|
self._db_path = db_path
|
|
self._config = config
|
|
self._context = context
|
|
self._rag = rag
|
|
self._lock = lock
|
|
self._search_results = []
|
|
self._doc_items = {}
|
|
self._doc_chunk_index = {}
|
|
self._items_jsonl_cache = {}
|
|
self._toc_json_cache = {}
|
|
self._repl = None
|
|
self._vfs = None
|
|
self._loop = None
|
|
|
|
@asynccontextmanager
|
|
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
|
|
"""Yield the shared connection (serialized by the lock), or an ephemeral
|
|
read-only one. The lock guards the whole block so a read's awaits cannot
|
|
interleave with another task's operation on the same connection."""
|
|
if self._rag is not None:
|
|
if self._lock is not None:
|
|
async with self._lock:
|
|
yield self._rag
|
|
else:
|
|
yield self._rag
|
|
return
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag:
|
|
yield rag
|
|
|
|
def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
|
|
"""Run a coroutine on the execute() loop from a synchronous callback.
|
|
|
|
Called from Monty's worker thread while ``feed_run_async`` leaves the
|
|
loop free, so scheduling onto it and blocking for the result is safe.
|
|
"""
|
|
assert self._loop is not None, (
|
|
"VFS reads happen during execute(); the loop must be captured first."
|
|
)
|
|
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
|
|
|
|
def close(self) -> None:
|
|
"""Retained for API compatibility; the sandbox owns no resources."""
|
|
return
|
|
|
|
def _build_external_functions(self) -> dict[str, Any]:
|
|
"""Build async external functions for the Monty interpreter."""
|
|
context = self._context
|
|
|
|
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
|
|
# Picture bytes are deliberately not attached to in-code search
|
|
# results: the Monty interpreter has no PIL/base64/hashlib, so the
|
|
# agent's Python can't do anything with them. The driving model
|
|
# gets figures through the top-level `search` tool when the
|
|
# question is visual; in-code search is for structural work.
|
|
async with self._connection() as rag:
|
|
results = await rag.search(query, limit=limit, filter=context.filter)
|
|
expanded = await rag.expand_context(results)
|
|
self._search_results.extend(expanded)
|
|
out: list[dict[str, Any]] = []
|
|
for r in expanded:
|
|
picture_refs = [
|
|
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
|
|
]
|
|
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]]:
|
|
async with self._connection() as rag:
|
|
docs = await rag.list_documents(filter=context.filter)
|
|
return [
|
|
{
|
|
"id": d.id,
|
|
"title": d.title,
|
|
"uri": d.uri,
|
|
"created_at": str(d.created_at),
|
|
}
|
|
for d in docs
|
|
]
|
|
|
|
return {
|
|
"search": search,
|
|
"list_documents": list_documents,
|
|
}
|
|
|
|
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, bulk-cached)
|
|
- toc.json: CallbackFile (lazy, bulk-cached)
|
|
"""
|
|
files: list[MemoryFile | CallbackFile] = []
|
|
|
|
def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
|
|
raise PermissionError(f"Document files are read-only: {_path}")
|
|
|
|
async with self._connection() as rag:
|
|
docs = await rag.list_documents(filter=self._context.filter)
|
|
|
|
doc_titles = {doc.id: doc.title for doc in docs if doc.id}
|
|
|
|
sandbox = self
|
|
|
|
def _get_items(did: str) -> list[DocumentItem]:
|
|
"""Fetch items for one doc, cached on the sandbox."""
|
|
cached = sandbox._doc_items.get(did)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
async def _fetch() -> list[DocumentItem]:
|
|
async with sandbox._connection() as rag:
|
|
return await rag.document_item_repository.get_all_items(did)
|
|
|
|
items = sandbox._run_on_loop(_fetch())
|
|
sandbox._doc_items[did] = items
|
|
return items
|
|
|
|
def _get_chunk_index(did: str) -> dict[str, list[str]]:
|
|
"""Fetch the self_ref → chunk_ids index for one doc, cached."""
|
|
cached = sandbox._doc_chunk_index.get(did)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
async def _fetch() -> dict[str, list[str]]:
|
|
async with sandbox._connection() as rag:
|
|
index = (
|
|
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
|
|
[did]
|
|
)
|
|
)
|
|
return index.get(did, {})
|
|
|
|
chunk_index = sandbox._run_on_loop(_fetch())
|
|
sandbox._doc_chunk_index[did] = chunk_index
|
|
return chunk_index
|
|
|
|
def _make_items_reader(
|
|
did: str,
|
|
) -> Callable[["PurePosixPath"], str]:
|
|
def read_items(_path: "PurePosixPath") -> str:
|
|
cached = sandbox._items_jsonl_cache.get(did)
|
|
if cached is not None:
|
|
return cached
|
|
items = _get_items(did)
|
|
chunk_index = _get_chunk_index(did)
|
|
jsonl = "\n".join(
|
|
json.dumps(
|
|
{
|
|
"self_ref": item.self_ref,
|
|
"label": item.label,
|
|
"text": item.text,
|
|
"page_numbers": item.page_numbers,
|
|
"heading_level": item.heading_level,
|
|
"chunk_ids": chunk_index.get(item.self_ref, []),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
for item in items
|
|
)
|
|
sandbox._items_jsonl_cache[did] = jsonl
|
|
return jsonl
|
|
|
|
return read_items
|
|
|
|
def _make_toc_reader(
|
|
did: str,
|
|
) -> Callable[["PurePosixPath"], str]:
|
|
def read_toc(_path: "PurePosixPath") -> str:
|
|
cached = sandbox._toc_json_cache.get(did)
|
|
if cached is not None:
|
|
return cached
|
|
items = _get_items(did)
|
|
chunk_index = _get_chunk_index(did)
|
|
toc = json.dumps(
|
|
{
|
|
"doc_id": did,
|
|
"title": doc_titles.get(did),
|
|
"tree": _build_toc(items, chunk_index),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
sandbox._toc_json_cache[did] = toc
|
|
return toc
|
|
|
|
return read_toc
|
|
|
|
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:
|
|
async with sandbox._connection() as rag:
|
|
content = await rag.document_repository.get_content(did)
|
|
return content or ""
|
|
|
|
return sandbox._run_on_loop(_fetch())
|
|
|
|
return read_content
|
|
|
|
files.append(
|
|
CallbackFile(
|
|
f"{doc_dir}/content.txt",
|
|
read=_make_content_reader(doc_id),
|
|
write=_deny_write,
|
|
)
|
|
)
|
|
files.append(
|
|
CallbackFile(
|
|
f"{doc_dir}/items.jsonl",
|
|
read=_make_items_reader(doc_id),
|
|
write=_deny_write,
|
|
)
|
|
)
|
|
# HAIKU_RAG_DISABLE_TOC is an evaluation-time toggle for measuring
|
|
# whether toc.json's outline view earns its place in the VFS.
|
|
# Production callers should leave it unset.
|
|
if not os.environ.get("HAIKU_RAG_DISABLE_TOC"):
|
|
files.append(
|
|
CallbackFile(
|
|
f"{doc_dir}/toc.json",
|
|
read=_make_toc_reader(doc_id),
|
|
write=_deny_write,
|
|
)
|
|
)
|
|
|
|
return OSAccess(files)
|
|
|
|
async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]:
|
|
"""Initialize the REPL session and VFS on first use."""
|
|
if self._repl is None:
|
|
self._vfs = await self._build_vfs()
|
|
self._repl = MontyRepl(
|
|
limits={
|
|
"max_duration_secs": self._config.analysis.code_timeout,
|
|
},
|
|
)
|
|
assert self._repl is not None and self._vfs is not None
|
|
return self._repl, self._vfs
|
|
|
|
async def execute(self, code: str) -> SandboxResult:
|
|
"""Execute Python code in the Monty REPL.
|
|
|
|
Variables persist across calls within the same Sandbox instance.
|
|
"""
|
|
# Monty's synchronous file callbacks bridge DB reads back to this loop.
|
|
self._loop = asyncio.get_running_loop()
|
|
repl, vfs = await self._ensure_initialized()
|
|
external_fns = self._build_external_functions()
|
|
|
|
stdout_lines: list[str] = []
|
|
|
|
def print_callback(_stream: Literal["stdout"], text: str) -> None:
|
|
stdout_lines.append(text)
|
|
|
|
max_chars = self._config.analysis.max_output_chars
|
|
|
|
try:
|
|
output = await repl.feed_run_async(
|
|
code,
|
|
external_functions=external_fns,
|
|
print_callback=print_callback,
|
|
os=vfs,
|
|
)
|
|
except (
|
|
pydantic_monty.MontySyntaxError,
|
|
pydantic_monty.MontyRuntimeError,
|
|
) as e:
|
|
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)
|
|
|
|
stdout = "".join(stdout_lines)
|
|
if output is not None:
|
|
stdout_with_output = f"{stdout}{output}" if stdout else str(output)
|
|
else:
|
|
stdout_with_output = stdout
|
|
|
|
if len(stdout_with_output) > max_chars:
|
|
stdout_with_output = (
|
|
stdout_with_output[:max_chars] + "\n... (output truncated)"
|
|
)
|
|
|
|
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)
|