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 import json
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty 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.agents.analysis.dependencies import AnalysisContext
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
@ -15,8 +16,6 @@ from haiku.rag.store.models.chunk import SearchResult
if TYPE_CHECKING: if TYPE_CHECKING:
from pathlib import PurePosixPath from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@dataclass @dataclass
class SandboxResult: class SandboxResult:
@ -41,35 +40,46 @@ class Sandbox:
and resolved asynchronously on the host. and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``. Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
sandbox = Sandbox(client, config, context) The interpreter uses a REPL session — variables persist across
result = await sandbox.execute("print('hello')") ``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 _config: AppConfig
_context: AnalysisContext _context: AnalysisContext
_search_results: "list[SearchResult]" _search_results: "list[SearchResult]"
_repl: MontyRepl | None
_vfs: OSAccess | None
def __init__( def __init__(
self, self,
client: "HaikuRAG", db_path: Path,
config: AppConfig, config: AppConfig,
context: AnalysisContext, context: AnalysisContext,
): ):
self._client = client self._db_path = db_path
self._config = config self._config = config
self._context = context self._context = context
self._search_results = [] self._search_results = []
self._repl = None
self._vfs = None
def _build_external_functions(self) -> dict[str, Any]: def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter.""" """Build async external functions for the Monty interpreter."""
client = self._client db_path = self._db_path
config = self._config config = self._config
context = self._context context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]: async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
results = await client.search(query, limit=limit, filter=context.filter) from haiku.rag.client import HaikuRAG
expanded = await client.expand_context(results)
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) self._search_results.extend(expanded)
return [ return [
{ {
@ -88,7 +98,10 @@ class Sandbox:
] ]
async def list_documents() -> list[dict[str, Any]]: 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 [ return [
{ {
"id": d.id, "id": d.id,
@ -123,10 +136,14 @@ class Sandbox:
- content.txt: CallbackFile (lazy, can be large) - content.txt: CallbackFile (lazy, can be large)
- items.jsonl: 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] = [] 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: for doc in docs:
if not doc.id: if not doc.id:
@ -150,17 +167,21 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]: ) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str: def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str: async def _fetch() -> str:
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import escape_sql_string from haiku.rag.utils import escape_sql_string
safe_id = escape_sql_string(did) async with HaikuRAG(
rows = list( db_path, config=config, read_only=True
client.store.documents_table.search() ) as rag:
.select(["content"]) safe_id = escape_sql_string(did)
.where(f"id = '{safe_id}'") rows = list(
.limit(1) rag.store.documents_table.search()
.to_list() .select(["content"])
) .where(f"id = '{safe_id}'")
return rows[0]["content"] if rows else "" .limit(1)
.to_list()
)
return rows[0]["content"] if rows else ""
return _run_async(_fetch()) return _run_async(_fetch())
@ -171,11 +192,16 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]: ) -> Callable[["PurePosixPath"], str]:
def read_items(_path: "PurePosixPath") -> str: def read_items(_path: "PurePosixPath") -> str:
async def _fetch() -> str: async def _fetch() -> str:
items = ( from haiku.rag.client import HaikuRAG
await client.document_item_repository.get_items_in_range(
did, 0, 999999 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 = [] lines = []
for item in items: for item in items:
lines.append( lines.append(
@ -213,37 +239,44 @@ class Sandbox:
return OSAccess(files) return OSAccess(files)
async def execute(self, code: str) -> SandboxResult: async def _ensure_initialized(self) -> None:
"""Execute Python code in the Monty interpreter.""" """Initialize the REPL session and VFS on first use."""
external_fns = self._build_external_functions() if self._repl is None:
vfs = await self._build_vfs() self._vfs = await self._build_vfs()
self._repl = MontyRepl(
input_names: list[str] = [] limits={
inputs: dict[str, Any] | None = None "max_duration_secs": self._config.analysis.code_timeout,
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,
) )
except ( if self._context.documents:
pydantic_monty.MontySyntaxError, await pydantic_monty.run_repl_async(
pydantic_monty.MontyRuntimeError, self._repl,
) as e: "pass",
return SandboxResult(stdout="", stderr=str(e), success=False) 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] = [] stdout_lines: list[str] = []
@ -251,20 +284,19 @@ class Sandbox:
stdout_lines.append(text) stdout_lines.append(text)
max_chars = self._config.analysis.max_output_chars max_chars = self._config.analysis.max_output_chars
limits: pydantic_monty.ResourceLimits = {
"max_duration_secs": self._config.analysis.code_timeout,
}
try: try:
output = await pydantic_monty.run_monty_async( output = await pydantic_monty.run_repl_async(
monty, self._repl,
inputs=inputs, code,
external_functions=external_fns, external_functions=external_fns,
limits=limits,
print_callback=print_callback, 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) stdout = "".join(stdout_lines)
if len(stdout) > max_chars: if len(stdout) > max_chars:
stdout = stdout[:max_chars] + "\n... (output truncated)" 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 context.documents = loaded_docs if loaded_docs else None
sandbox = Sandbox( sandbox = Sandbox(
client=self, db_path=self.store.db_path,
config=self._config, config=self._config,
context=context, 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 tools["get_document"] = get_document
if "execute_code" in tool_names: 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: async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter. """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 and a virtual filesystem at /documents/ with document content and
structure (metadata.json, content.txt, items.jsonl per document). structure (metadata.json, content.txt, items.jsonl per document).
Use print() to output results. Each call runs in a fresh Use print() to output results. Variables persist between calls.
interpreter — variables do not persist between calls.
Args: Args:
code: Python code to execute. code: Python code to execute.
""" """
from haiku.rag.agents.analysis.dependencies import AnalysisContext from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox 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) state = _get_state(ctx, state_type)
doc_filter = state.document_filter if state else None if state and sandbox._search_results:
context = AnalysisContext(filter=doc_filter) existing = state.searches.get("_sandbox", [])
seen = {r.chunk_id for r in existing}
async with HaikuRAG(db_path, config=config, read_only=True) as rag: for sr in sandbox._search_results:
sandbox = Sandbox(client=rag, config=config, context=context) if sr.chunk_id not in seen:
result = await sandbox.execute(code) existing.append(sr)
seen.add(sr.chunk_id)
if state and sandbox._search_results: state.searches["_sandbox"] = existing
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: if state:
state.executions.append( 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 ## Tools
### execute_code ### 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`): 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
@ -71,7 +71,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
## Important ## 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 - 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 - 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, llm)

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

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