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

use MontyRepl for persistent variables across execute_code calls

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-20 10:47:50 +03:00
parent fa87cf79c5
commit 4d75943df5
No known key found for this signature in database
6 changed files with 186 additions and 146 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3,10 +3,11 @@ import concurrent.futures
import json
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from pydantic_monty import CallbackFile, MemoryFile, OSAccess
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.config.models import AppConfig
@ -15,8 +16,6 @@ from haiku.rag.store.models.chunk import SearchResult
if TYPE_CHECKING:
from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@dataclass
class SandboxResult:
@ -41,35 +40,46 @@ class Sandbox:
and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
sandbox = Sandbox(client, config, context)
result = await sandbox.execute("print('hello')")
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
"""
_client: "HaikuRAG"
_db_path: Path
_config: AppConfig
_context: AnalysisContext
_search_results: "list[SearchResult]"
_repl: MontyRepl | None
_vfs: OSAccess | None
def __init__(
self,
client: "HaikuRAG",
db_path: Path,
config: AppConfig,
context: AnalysisContext,
):
self._client = client
self._db_path = db_path
self._config = config
self._context = context
self._search_results = []
self._repl = None
self._vfs = None
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
client = self._client
db_path = self._db_path
config = self._config
context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
results = await client.search(query, limit=limit, filter=context.filter)
expanded = await client.expand_context(results)
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(query, limit=limit, filter=context.filter)
expanded = await rag.expand_context(results)
self._search_results.extend(expanded)
return [
{
@ -88,7 +98,10 @@ class Sandbox:
]
async def list_documents() -> list[dict[str, Any]]:
docs = await client.list_documents(filter=context.filter)
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
docs = await rag.list_documents(filter=context.filter)
return [
{
"id": d.id,
@ -123,10 +136,14 @@ class Sandbox:
- content.txt: CallbackFile (lazy, can be large)
- items.jsonl: CallbackFile (lazy, can be large)
"""
client = self._client
from haiku.rag.client import HaikuRAG
db_path = self._db_path
config = self._config
files: list[MemoryFile | CallbackFile] = []
docs = await client.list_documents(filter=self._context.filter)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
docs = await rag.list_documents(filter=self._context.filter)
for doc in docs:
if not doc.id:
@ -150,17 +167,21 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import escape_sql_string
safe_id = escape_sql_string(did)
rows = list(
client.store.documents_table.search()
.select(["content"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
return rows[0]["content"] if rows else ""
async with HaikuRAG(
db_path, config=config, read_only=True
) as rag:
safe_id = escape_sql_string(did)
rows = list(
rag.store.documents_table.search()
.select(["content"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
return rows[0]["content"] if rows else ""
return _run_async(_fetch())
@ -171,11 +192,16 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]:
def read_items(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
items = (
await client.document_item_repository.get_items_in_range(
did, 0, 999999
from haiku.rag.client import HaikuRAG
async with HaikuRAG(
db_path, config=config, read_only=True
) as rag:
items = (
await rag.document_item_repository.get_items_in_range(
did, 0, 999999
)
)
)
lines = []
for item in items:
lines.append(
@ -213,37 +239,44 @@ class Sandbox:
return OSAccess(files)
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty interpreter."""
external_fns = self._build_external_functions()
vfs = await self._build_vfs()
input_names: list[str] = []
inputs: dict[str, Any] | None = None
if self._context.documents:
input_names.append("documents")
inputs = {
"documents": [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"content": d.content,
}
for d in self._context.documents
]
}
try:
monty = pydantic_monty.Monty(
code,
inputs=input_names,
async def _ensure_initialized(self) -> None:
"""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,
},
)
except (
pydantic_monty.MontySyntaxError,
pydantic_monty.MontyRuntimeError,
) as e:
return SandboxResult(stdout="", stderr=str(e), success=False)
if self._context.documents:
await pydantic_monty.run_repl_async(
self._repl,
"pass",
inputs={
"documents": [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"content": d.content,
}
for d in self._context.documents
]
},
external_functions=self._build_external_functions(),
os=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.
"""
await self._ensure_initialized()
assert self._repl is not None
assert self._vfs is not None
external_fns = self._build_external_functions()
stdout_lines: list[str] = []
@ -251,20 +284,19 @@ class Sandbox:
stdout_lines.append(text)
max_chars = self._config.analysis.max_output_chars
limits: pydantic_monty.ResourceLimits = {
"max_duration_secs": self._config.analysis.code_timeout,
}
try:
output = await pydantic_monty.run_monty_async(
monty,
inputs=inputs,
output = await pydantic_monty.run_repl_async(
self._repl,
code,
external_functions=external_fns,
limits=limits,
print_callback=print_callback,
os=vfs,
os=self._vfs,
)
except pydantic_monty.MontyRuntimeError as e:
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)"

2
haiku_rag_slim/haiku/rag/client.py
View file

@ -1230,7 +1230,7 @@ class HaikuRAG:
context.documents = loaded_docs if loaded_docs else None
sandbox = Sandbox(
client=self,
db_path=self.store.db_path,
config=self._config,
context=context,
)

39
haiku_rag_slim/haiku/rag/skills/_tools.py
View file

@ -240,6 +240,7 @@ def create_skill_tools(
tools["get_document"] = get_document
if "execute_code" in tool_names:
_sandbox: list[Any] = [] # mutable container for closure; holds [Sandbox] or []
async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter.
@ -248,32 +249,34 @@ def create_skill_tools(
and a virtual filesystem at /documents/ with document content and
structure (metadata.json, content.txt, items.jsonl per document).
Use print() to output results. Each call runs in a fresh
interpreter — variables do not persist between calls.
Use print() to output results. Variables persist between calls.
Args:
code: Python code to execute.
"""
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
if not _sandbox:
state = _get_state(ctx, state_type)
doc_filter = state.document_filter if state else None
context = AnalysisContext(filter=doc_filter)
_sandbox.append(
Sandbox(db_path=db_path, config=config, context=context)
)
sandbox = _sandbox[0]
result = await sandbox.execute(code)
state = _get_state(ctx, state_type)
doc_filter = state.document_filter if state else None
context = AnalysisContext(filter=doc_filter)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
sandbox = Sandbox(client=rag, config=config, context=context)
result = await sandbox.execute(code)
if state and sandbox._search_results:
existing = state.searches.get("_sandbox", [])
seen = {r.chunk_id for r in existing}
for sr in sandbox._search_results:
if sr.chunk_id not in seen:
existing.append(sr)
seen.add(sr.chunk_id)
state.searches["_sandbox"] = existing
if state and sandbox._search_results:
existing = state.searches.get("_sandbox", [])
seen = {r.chunk_id for r in existing}
for sr in sandbox._search_results:
if sr.chunk_id not in seen:
existing.append(sr)
seen.add(sr.chunk_id)
state.searches["_sandbox"] = existing
if state:
state.executions.append(

4
haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
View file

@ -15,7 +15,7 @@ You solve complex analytical questions by writing and executing Python code agai
## Tools
### execute_code
Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — variables do not persist between calls. Use `print()` to output results.
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
@ -71,7 +71,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
## Important
- Each `execute_code` call runs in a fresh interpreter (no persistent variables between calls)
- 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)

9
tests/agents/analysis/conftest.py
View file

@ -14,8 +14,9 @@ async def empty_client(temp_db_path):
@pytest.fixture
async def sandbox(empty_client):
async def sandbox(temp_db_path):
"""Create a Monty sandbox for testing."""
config = AppConfig()
context = AnalysisContext()
return Sandbox(client=empty_client, config=config, context=context)
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
context = AnalysisContext()
return Sandbox(db_path=temp_db_path, config=config, context=context)

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

@ -99,7 +99,7 @@ class TestSandboxListDocuments:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"docs = await list_documents()\n"
"print(len(docs))\n"
@ -126,7 +126,7 @@ class TestSandboxSearch:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('fox', limit=5)\n"
"print(len(results))\n"
@ -149,7 +149,7 @@ class TestSandboxSearch:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('fox', limit=1)\n"
"r = results[0]\n"
@ -175,7 +175,7 @@ class TestSandboxSearch:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('fox', limit=1)\n"
"print(type(results[0]['content']).__name__)\n"
@ -242,29 +242,31 @@ class TestSandboxOutputTruncation:
"""Test output truncation behavior."""
@pytest.mark.asyncio
async def test_truncate_stdout_on_runtime_error(self, empty_client):
async def test_truncate_stdout_on_runtime_error(self, temp_db_path):
"""Test stdout is truncated when a runtime error occurs after large output."""
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('a' * 100)\nx = 1/0")
assert not result.success
assert "ZeroDivisionError" in result.stderr
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute("print('a' * 100)\nx = 1/0")
assert not result.success
assert "ZeroDivisionError" in result.stderr
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
@pytest.mark.asyncio
async def test_truncate_successful_output(self, empty_client):
async def test_truncate_successful_output(self, temp_db_path):
"""Test output is truncated on successful execution with large output."""
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('b' * 100)")
assert result.success
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute("print('b' * 100)")
assert result.success
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
class TestSandboxVFS:
@ -293,7 +295,7 @@ class TestSandboxVFS:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"dirs = list(Path('/documents').iterdir())\n"
@ -317,7 +319,7 @@ class TestSandboxVFS:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
@ -342,7 +344,7 @@ class TestSandboxVFS:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
f"content = Path('/documents/{doc.id}/content.txt').read_text()\n"
@ -364,7 +366,7 @@ class TestSandboxVFS:
)
context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
@ -399,7 +401,7 @@ class TestSandboxVFS:
)
context = AnalysisContext(filter="uri LIKE 'public://%'")
sb = Sandbox(client=client, config=config, context=context)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
@ -425,24 +427,25 @@ class TestSandboxPreloadedDocuments:
assert "NameError" in result.stderr
@pytest.mark.asyncio
async def test_documents_variable_available_with_preload(self, empty_client):
async def test_documents_variable_available_with_preload(self, temp_db_path):
"""documents variable is available when context.documents is set."""
config = AppConfig()
docs = [
Document(id="1", content="Content A", title="Doc A", uri="a://1"),
Document(id="2", content="Content B", title="Doc B", uri="b://2"),
]
context = AnalysisContext(documents=docs)
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute(
"print(len(documents))\n"
"print(documents[0]['title'])\n"
"print(documents[1]['title'])"
)
assert result.success
assert "2" in result.stdout
assert "Doc A" in result.stdout
assert "Doc B" in result.stdout
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
docs = [
Document(id="1", content="Content A", title="Doc A", uri="a://1"),
Document(id="2", content="Content B", title="Doc B", uri="b://2"),
]
context = AnalysisContext(documents=docs)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"print(len(documents))\n"
"print(documents[0]['title'])\n"
"print(documents[1]['title'])"
)
assert result.success
assert "2" in result.stdout
assert "Doc A" in result.stdout
assert "Doc B" in result.stdout
class TestSandboxLLM:
@ -450,14 +453,15 @@ class TestSandboxLLM:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_llm_function(self, allow_model_requests, empty_client):
async def test_llm_function(self, allow_model_requests, temp_db_path):
"""Test llm() calls the model and returns a string."""
config = AppConfig()
context = AnalysisContext()
sb = Sandbox(client=empty_client, 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
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