From 4d75943df5ae80cb3bb537d946e31e69d46ca97f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 20 Apr 2026 10:47:50 +0300 Subject: [PATCH] use MontyRepl for persistent variables across execute_code calls --- .../haiku/rag/agents/analysis/sandbox.py | 164 +++++++++++------- haiku_rag_slim/haiku/rag/client.py | 2 +- haiku_rag_slim/haiku/rag/skills/_tools.py | 39 +++-- .../haiku/rag/skills/rag-analysis/SKILL.md | 4 +- tests/agents/analysis/conftest.py | 9 +- tests/agents/analysis/test_sandbox.py | 114 ++++++------ 6 files changed, 186 insertions(+), 146 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index c39ec1b3..635e9caa 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -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)" diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index c8473dd3..d83ce3fc 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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, ) diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 8e7f9270..22c17b8b 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -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( diff --git a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md index 31be509d..784e1e39 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md +++ b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md @@ -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) diff --git a/tests/agents/analysis/conftest.py b/tests/agents/analysis/conftest.py index 55810a53..ca2f3712 100644 --- a/tests/agents/analysis/conftest.py +++ b/tests/agents/analysis/conftest.py @@ -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) diff --git a/tests/agents/analysis/test_sandbox.py b/tests/agents/analysis/test_sandbox.py index 4b846f6c..26bc4679 100644 --- a/tests/agents/analysis/test_sandbox.py +++ b/tests/agents/analysis/test_sandbox.py @@ -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