Fix analysis tool to use docker sandbox

This commit is contained in:
Yiorgis Gozadinos 2026-02-05 22:12:09 +01:00
parent 3ed9cbb7d3
commit 84b320765a
No known key found for this signature in database
4 changed files with 28 additions and 40 deletions

View file

@ -3,7 +3,7 @@ from pydantic_ai import FunctionToolset
from haiku.rag.agents.rlm.agent import create_rlm_agent from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution from haiku.rag.agents.rlm.docker_sandbox import DockerSandbox
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext from haiku.rag.tools.context import ToolContext
@ -21,10 +21,10 @@ ANALYSIS_NAMESPACE = "haiku.rag.analysis"
class AnalysisState(BaseModel): class AnalysisState(BaseModel):
"""State for analysis toolset. """State for analysis toolset.
Tracks code executions across tool invocations. Tracks programs produced across tool invocations.
""" """
code_executions: list[CodeExecution] = [] programs: list[str] = []
def create_analysis_toolset( def create_analysis_toolset(
@ -85,28 +85,32 @@ def create_analysis_toolset(
combine_filters(base_filter, session_filter), doc_filter combine_filters(base_filter, session_filter), doc_filter
) )
# Create RLM context and deps # Create RLM context
rlm_context = RLMContext(filter=effective_filter) rlm_context = RLMContext(filter=effective_filter)
deps = RLMDeps(
# Run RLM agent with Docker sandbox
async with DockerSandbox(
client=client, client=client,
config=config, config=config.rlm,
context=rlm_context, context=rlm_context,
) image=config.rlm.docker_image,
) as sandbox:
deps = RLMDeps(
sandbox=sandbox,
context=rlm_context,
)
# Run RLM agent rlm_agent = create_rlm_agent(config)
rlm_agent = create_rlm_agent(config) result = await rlm_agent.run(task, deps=deps)
result = await rlm_agent.run(task, deps=deps)
# Track code executions in state program = result.output.program
code_executions = rlm_context.code_executions if state is not None and program:
if state is not None: state.programs.append(program)
state.code_executions.extend(code_executions)
return AnalysisResult( return AnalysisResult(
answer=result.output.answer, answer=result.output.answer,
code_executed=len(code_executions) > 0, code_executed=bool(program),
execution_count=len(code_executions), )
)
toolset = FunctionToolset() toolset = FunctionToolset()
toolset.add_function(analyze, name=tool_name) toolset.add_function(analyze, name=tool_name)

View file

@ -35,7 +35,3 @@ class AnalysisResult(BaseModel):
default=True, default=True,
description="Whether code was executed to produce this answer", description="Whether code was executed to produce this answer",
) )
execution_count: int = Field(
default=0,
description="Number of code executions performed",
)

View file

@ -12,28 +12,19 @@ class TestAnalysisState:
"""Tests for AnalysisState model.""" """Tests for AnalysisState model."""
def test_analysis_state_defaults(self): def test_analysis_state_defaults(self):
"""AnalysisState initializes with empty code_executions.""" """AnalysisState initializes with empty programs."""
state = AnalysisState() state = AnalysisState()
assert state.code_executions == [] assert state.programs == []
def test_analysis_state_serialization(self): def test_analysis_state_serialization(self):
"""AnalysisState serializes and deserializes correctly.""" """AnalysisState serializes and deserializes correctly."""
from haiku.rag.agents.rlm.models import CodeExecution
state = AnalysisState() state = AnalysisState()
state.code_executions.append( state.programs.append("print('hello')")
CodeExecution(
code="print('hello')",
stdout="hello\n",
stderr="",
success=True,
)
)
data = state.model_dump() data = state.model_dump()
restored = AnalysisState.model_validate(data) restored = AnalysisState.model_validate(data)
assert len(restored.code_executions) == 1 assert len(restored.programs) == 1
assert restored.code_executions[0].code == "print('hello')" assert restored.programs[0] == "print('hello')"
class TestAnalysisToolset: class TestAnalysisToolset:

View file

@ -84,7 +84,6 @@ def test_analysis_result_defaults():
"""Test AnalysisResult has sensible defaults.""" """Test AnalysisResult has sensible defaults."""
result = AnalysisResult(answer="The result is 42") result = AnalysisResult(answer="The result is 42")
assert result.code_executed is True assert result.code_executed is True
assert result.execution_count == 0
def test_analysis_result_with_values(): def test_analysis_result_with_values():
@ -92,8 +91,6 @@ def test_analysis_result_with_values():
result = AnalysisResult( result = AnalysisResult(
answer="The result is 42", answer="The result is 42",
code_executed=True, code_executed=True,
execution_count=3,
) )
assert result.answer == "The result is 42" assert result.answer == "The result is 42"
assert result.code_executed is True assert result.code_executed is True
assert result.execution_count == 3