From 84b320765a342c19a331fc99e38c7eacc54fb50b Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 5 Feb 2026 22:12:09 +0100
Subject: [PATCH] Fix analysis tool to use docker sandbox

---
 haiku_rag_slim/haiku/rag/tools/analysis.py | 42 ++++++++++++----------
 haiku_rag_slim/haiku/rag/tools/models.py   |  4 ---
 tests/tools/test_analysis.py               | 19 +++-------
 tests/tools/test_models.py                 |  3 --
 4 files changed, 28 insertions(+), 40 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/tools/analysis.py b/haiku_rag_slim/haiku/rag/tools/analysis.py
index e95386c8..8e987bdd 100644
--- a/haiku_rag_slim/haiku/rag/tools/analysis.py
+++ b/haiku_rag_slim/haiku/rag/tools/analysis.py
@@ -3,7 +3,7 @@ from pydantic_ai import FunctionToolset
 
 from haiku.rag.agents.rlm.agent import create_rlm_agent
 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.config.models import AppConfig
 from haiku.rag.tools.context import ToolContext
@@ -21,10 +21,10 @@ ANALYSIS_NAMESPACE = "haiku.rag.analysis"
 class AnalysisState(BaseModel):
     """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(
@@ -85,28 +85,32 @@ def create_analysis_toolset(
             combine_filters(base_filter, session_filter), doc_filter
         )
 
-        # Create RLM context and deps
+        # Create RLM context
         rlm_context = RLMContext(filter=effective_filter)
-        deps = RLMDeps(
+
+        # Run RLM agent with Docker sandbox
+        async with DockerSandbox(
             client=client,
-            config=config,
+            config=config.rlm,
             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)
-        result = await rlm_agent.run(task, deps=deps)
+            rlm_agent = create_rlm_agent(config)
+            result = await rlm_agent.run(task, deps=deps)
 
-        # Track code executions in state
-        code_executions = rlm_context.code_executions
-        if state is not None:
-            state.code_executions.extend(code_executions)
+            program = result.output.program
+            if state is not None and program:
+                state.programs.append(program)
 
-        return AnalysisResult(
-            answer=result.output.answer,
-            code_executed=len(code_executions) > 0,
-            execution_count=len(code_executions),
-        )
+            return AnalysisResult(
+                answer=result.output.answer,
+                code_executed=bool(program),
+            )
 
     toolset = FunctionToolset()
     toolset.add_function(analyze, name=tool_name)
diff --git a/haiku_rag_slim/haiku/rag/tools/models.py b/haiku_rag_slim/haiku/rag/tools/models.py
index d911a338..6d69b5df 100644
--- a/haiku_rag_slim/haiku/rag/tools/models.py
+++ b/haiku_rag_slim/haiku/rag/tools/models.py
@@ -35,7 +35,3 @@ class AnalysisResult(BaseModel):
         default=True,
         description="Whether code was executed to produce this answer",
     )
-    execution_count: int = Field(
-        default=0,
-        description="Number of code executions performed",
-    )
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index f4fac3e0..b01c8832 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -12,28 +12,19 @@ class TestAnalysisState:
     """Tests for AnalysisState model."""
 
     def test_analysis_state_defaults(self):
-        """AnalysisState initializes with empty code_executions."""
+        """AnalysisState initializes with empty programs."""
         state = AnalysisState()
-        assert state.code_executions == []
+        assert state.programs == []
 
     def test_analysis_state_serialization(self):
         """AnalysisState serializes and deserializes correctly."""
-        from haiku.rag.agents.rlm.models import CodeExecution
-
         state = AnalysisState()
-        state.code_executions.append(
-            CodeExecution(
-                code="print('hello')",
-                stdout="hello\n",
-                stderr="",
-                success=True,
-            )
-        )
+        state.programs.append("print('hello')")
 
         data = state.model_dump()
         restored = AnalysisState.model_validate(data)
-        assert len(restored.code_executions) == 1
-        assert restored.code_executions[0].code == "print('hello')"
+        assert len(restored.programs) == 1
+        assert restored.programs[0] == "print('hello')"
 
 
 class TestAnalysisToolset:
diff --git a/tests/tools/test_models.py b/tests/tools/test_models.py
index 65a83a78..c01c8990 100644
--- a/tests/tools/test_models.py
+++ b/tests/tools/test_models.py
@@ -84,7 +84,6 @@ def test_analysis_result_defaults():
     """Test AnalysisResult has sensible defaults."""
     result = AnalysisResult(answer="The result is 42")
     assert result.code_executed is True
-    assert result.execution_count == 0
 
 
 def test_analysis_result_with_values():
@@ -92,8 +91,6 @@ def test_analysis_result_with_values():
     result = AnalysisResult(
         answer="The result is 42",
         code_executed=True,
-        execution_count=3,
     )
     assert result.answer == "The result is 42"
     assert result.code_executed is True
-    assert result.execution_count == 3