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

persist sandbox variables across execute_code calls within one invocation

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-22 13:30:38 +03:00
parent 4e9c02afc2
commit 4a9dd9b49a
No known key found for this signature in database
8 changed files with 172 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -5,6 +5,7 @@
- **Skills share a single `HaikuRAG` client per invocation** via the new `haiku.skills>=0.15.0` `lifespan` hook. The skill's sub-agent opens one read-only client on entry, all tool calls reuse it, and it closes on exit — replacing the old pattern of open/close around every `search` / `list_documents` / `get_document` call.
- **`max_searches` tracked on `RAGRunDeps.search_count`** instead of a module-level `ctx.run_id`-keyed dict. Eliminates a memory leak in long-running processes where old run ids were never evicted.
- **Analysis sandbox persists variables across `execute_code` calls within one invocation.** Re-enables the incremental-exploration workflow (search in one call, process results in the next). Each new skill invocation constructs a fresh `Sandbox` via the analysis lifespan, so there is no cross-invocation leak.
## [0.41.0] - 2026-04-20

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

@ -8,7 +8,7 @@ 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
@ -44,11 +44,12 @@ class Sandbox:
and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
Each ``execute()`` call runs in a fresh interpreter — variables do not
persist between calls.
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("print('hello')")
result = await sandbox.execute("x = await search('query')")
result = await sandbox.execute("print(x[0]['content'])") # x persists
"""
_db_path: Path
@ -56,6 +57,8 @@ class Sandbox:
_context: AnalysisContext
_search_results: "list[SearchResult]"
_items_cache: dict[str, str] | None
_repl: MontyRepl | None
_vfs: OSAccess | None
def __init__(
self,
@ -68,6 +71,8 @@ class Sandbox:
self._context = context
self._search_results = []
self._items_cache = None
self._repl = None
self._vfs = None
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
@ -245,37 +250,46 @@ 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) -> 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,
},
)
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,
)
repl = self._repl
vfs = self._vfs
if repl is None or vfs is None:
raise RuntimeError("Sandbox initialization failed")
return repl, vfs
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty REPL.
Variables persist across calls within the same Sandbox instance.
"""
repl, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions()
stdout_lines: list[str] = []
@ -283,20 +297,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(
repl,
code,
external_functions=external_fns,
limits=limits,
print_callback=print_callback,
os=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)"

21
haiku_rag_slim/haiku/rag/skills/_deps.py
View file

@ -34,3 +34,24 @@ def make_rag_lifespan(db_path: Path, config: AppConfig):
yield
return lifespan
def make_analysis_lifespan(db_path: Path, config: AppConfig):
@asynccontextmanager
async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]:
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
doc_filter = getattr(deps.state, "document_filter", None)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
deps.rag = rag
deps.search_count = 0
deps.sandbox = Sandbox(
db_path=db_path,
config=config,
context=AnalysisContext(filter=doc_filter),
)
yield
return lifespan

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

@ -227,27 +227,26 @@ def create_skill_tools(
tools["get_document"] = get_document
if "execute_code" in tool_names:
from haiku.rag.skills._deps import AnalysisRunDeps
async def execute_code(ctx: RunContext[RAGRunDeps], code: str) -> str:
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(), llm() functions
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
within the same skill invocation.
Args:
code: Python code to execute.
"""
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
state = _get_state(ctx, state_type)
doc_filter = state.document_filter if state else None
context = AnalysisContext(filter=doc_filter)
sandbox = Sandbox(db_path=db_path, config=config, context=context)
if ctx.deps is None or ctx.deps.sandbox is None:
raise RuntimeError(
"AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
)
sandbox = ctx.deps.sandbox
result = await sandbox.execute(code)
state = _get_state(ctx, state_type)

4
haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -60,7 +60,7 @@ def create_skill(
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
if config is None:
@ -95,5 +95,5 @@ def create_skill(
state_type=STATE_TYPE,
state_namespace=STATE_NAMESPACE,
deps_type=AnalysisRunDeps,
lifespan=make_rag_lifespan(db_path, config),
lifespan=make_analysis_lifespan(db_path, config),
)

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 — write self-contained code. 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
@ -93,7 +93,7 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
## Important
- Each `execute_code` call runs in a fresh interpreter — write self-contained code blocks
- 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)

22
tests/skills/conftest.py
View file

@ -12,13 +12,13 @@ from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
VECTOR_DIM = 2560
def _make_ctx(state=None, rag=None):
def _make_ctx(state=None, rag=None, sandbox=None):
"""Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
from haiku.rag.skills.analysis import AnalysisState
ctx = MagicMock(spec=RunContext)
if isinstance(state, AnalysisState):
ctx.deps = AnalysisRunDeps(state=state, rag=rag)
if isinstance(state, AnalysisState) or sandbox is not None:
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
else:
ctx.deps = RAGRunDeps(state=state, rag=rag)
return ctx
@ -80,3 +80,19 @@ async def rag_client(rag_db):
"""Yield an open read-only HaikuRAG client on the sample db."""
async with HaikuRAG(rag_db, read_only=True) as rag:
yield rag
@pytest.fixture
def sandbox_factory(rag_db, test_app_config):
"""Build Sandbox instances bound to the sample db, optionally with a doc filter."""
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
def _make(filter: str | None = None) -> Sandbox:
return Sandbox(
db_path=rag_db,
config=test_app_config,
context=AnalysisContext(filter=filter),
)
return _make

78
tests/skills/test_analysis.py
View file

@ -113,73 +113,75 @@ class TestDomainPreambleInAnalysisSkillInstructions:
class TestExecuteCodeTool:
async def test_execute_code_returns_output(self, rag_db):
async def test_execute_code_returns_output(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(ctx, code="print('hello')")
assert "hello" in result
async def test_execute_code_updates_state(self, rag_db):
async def test_execute_code_updates_state(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="print('hello')")
assert len(state.executions) == 1
assert state.executions[0].code == "print('hello')"
assert state.executions[0].success is True
assert "hello" in state.executions[0].stdout
async def test_execute_code_reports_errors(self, rag_db):
async def test_execute_code_reports_errors(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(ctx, code="x = 1/0")
assert "Error" in result
assert "ZeroDivisionError" in result
assert state.executions[0].success is False
async def test_execute_code_applies_document_filter(self, rag_db):
async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory(filter=state.document_filter))
result = await execute_code(
ctx, code="docs = await list_documents()\nprint(len(docs))"
)
assert "1" in result
async def test_execute_code_accumulates_search_results(self, rag_db):
async def test_execute_code_accumulates_search_results(
self, rag_db, sandbox_factory
):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(
ctx, code="results = await search('intelligence')\nprint(len(results))"
)
assert "_sandbox" in state.searches
assert len(state.searches["_sandbox"]) > 0
async def test_execute_code_vfs_write_denied(self, rag_db):
async def test_execute_code_vfs_write_denied(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(
ctx,
code=(
@ -193,21 +195,67 @@ class TestExecuteCodeTool:
assert "Error" in result
assert "read-only" in result
async def test_execute_code_variables_persist_within_invocation(
self, rag_db, sandbox_factory
):
"""Same sandbox across two calls → vars persist (one skill invocation)."""
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="x = 42")
result = await execute_code(ctx, code="print(x * 2)")
assert "84" in result
async def test_execute_code_isolated_across_invocations(
self, rag_db, sandbox_factory
):
"""Different Sandbox instances → no cross-invocation leak."""
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
ctx1 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
await execute_code(ctx1, code="secret = 'do not leak'")
ctx2 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
result = await execute_code(ctx2, code="print(secret)")
assert not result.startswith("do not leak")
assert "Error" in result or "NameError" in result
class TestAnalysisLifespan:
async def test_opens_one_client_per_invocation(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_rag_lifespan
async def test_opens_client_and_sandbox_per_invocation(self, rag_db):
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
config = AppConfig()
lifespan = make_rag_lifespan(rag_db, config)
lifespan = make_analysis_lifespan(rag_db, config)
deps = AnalysisRunDeps()
async with lifespan(deps):
assert deps.rag is not None
assert deps.rag.is_read_only
assert deps.search_count == 0
assert isinstance(deps.sandbox, Sandbox)
docs = await deps.rag.list_documents()
assert len(docs) == 2
async def test_lifespan_reads_document_filter_from_state(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
from haiku.rag.skills.analysis import AnalysisState
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
state = AnalysisState(document_filter="title = 'AI Overview'")
deps = AnalysisRunDeps(state=state)
async with lifespan(deps):
assert deps.sandbox is not None
assert deps.sandbox._context.filter == "title = 'AI Overview'"
async def test_skill_has_lifespan_and_deps_type(
self, test_app_config, temp_db_path
):