Fix analysis tool to use docker sandbox
This commit is contained in:
parent
3ed9cbb7d3
commit
84b320765a
4 changed files with 28 additions and 40 deletions
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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",
|
|
||||||
)
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue