Fix llm example tests

Nodes with errors are available in results, just as errors

Author: iantbutler01

asimov/graph/__init__.py

@@ -535,6 +535,8 @@ async def _run_task(self, task: Task) -> None:
                 for (i, _), result in zip(tasks, results):
                     node_name = parallel_group["nodes"][i]
 
+                    self.node_results[node_name] = result
+
                     if not self.is_success(result):
                         dependent_chain = self.get_dependent_chains(node_name)
                         failed_chains.add(dependent_chain)
@@ -549,8 +551,6 @@ async def _run_task(self, task: Task) -> None:
                         )
                         continue
 
-                    self.node_results[node_name] = result
-
                     self.execution_state.mark_executed(
                         self.execution_state.execution_index, i
                     )

examples/llm_agent.py

@@ -80,7 +80,11 @@ async def process(
                 timeout=30.0,
             )
 
-            response_content = json.loads(response)["steps"]
+            try:
+                loaded_response = json.loads(response)
+                response_content = loaded_response["steps"]
+            except json.JSONDecodeError as e:
+                return {"status": "error", "result": str(e)}
 
             print(f"{self.name}: Received plan from LLM")
         except asyncio.TimeoutError:

tests/test_agent.py

@@ -266,7 +266,7 @@ async def test_error_handling(simple_agent, mock_cache):
     assert "NodeB" in result["failed_chains"][0]
     assert "Executed A" in str(result["result"]["NodeA"])
     assert "Executed C" in str(result["result"]["NodeC"])
-    assert "NodeB" not in result["result"]
+    assert "NodeB" in result["result"]
     assert "NodeD" not in result["result"]
     assert error["status"] == "error" and "NodeB" in error["node"]
 
@@ -744,7 +744,7 @@ async def test_caching_results(simple_agent, mock_cache):
     result = await mock_cache.get_message(simple_agent.output_mailbox)
     assert "NodeA" in result["result"]
     assert "NodeB" in result["result"]
-    assert "NodeC" not in result["result"]
+    assert "NodeC" in result["result"]
 
 
 @pytest.mark.asyncio

tests/test_llm_agent.py

@@ -20,26 +20,21 @@
 from examples.llm_agent import LLMPlannerModule, LLMExecutorModule, LLMFlowControlModule
 
 
-class MockAnthropicResponse:
-    def __init__(self, content):
-        self.choices = [MagicMock(message=MagicMock(content=content))]
-
-
 class MockAnthropicClient:
     """Mock client for testing LLM interactions."""
 
     def __init__(self, responses=None):
         self.responses = responses or {}
-        self.default_response = MockAnthropicResponse("Default mock response")
+        self.default_response = {}
 
-    async def get_generation(self, messages: List[Any]) -> MockAnthropicResponse:
+    async def get_generation(self, messages: List[Any]) -> str:
         # Extract prompt from messages
         prompt = messages[-1].content if messages else ""
         # Return predefined response if available, otherwise default
         for key, response in self.responses.items():
             if key.lower() in prompt.lower():
-                return MockAnthropicResponse(response)
-        return self.default_response
+                return response
+        return json.dumps(self.default_response)
 
 
 @pytest.fixture(autouse=True)
@@ -50,6 +45,7 @@ def setup_env():
     if "ANTHROPIC_API_KEY" in os.environ:
         del os.environ["ANTHROPIC_API_KEY"]
 
+
 @pytest.fixture
 def mock_cache():
     return MockRedisCache()
@@ -59,7 +55,7 @@ def mock_cache():
 def mock_anthropic_client():
     return MockAnthropicClient(
         {
-            "Create a step-by-step plan": '[{"description": "Research AI agents", "requirements": "Access to documentation", "validation": "Comprehensive notes available"}, {"description": "Write introduction", "requirements": "Research notes", "validation": "Clear introduction exists"}]',
+            "Create a step-by-step plan": '{"steps": [{"description": "Research AI agents", "requirements": "Access to documentation", "validation": "Comprehensive notes available"}, {"description": "Write introduction", "requirements": "Research notes", "validation": "Clear introduction exists"}]}',
             "Execute this step": "Step executed successfully with the following results:\n1. Actions: Researched AI agents\n2. Outcome: Comprehensive notes created\n3. Output: 5 pages of detailed notes",
             "Evaluate if the step": "success - all validation criteria met",
             "Analyze the execution history": "Analysis complete. Decision: continue - execution is proceeding as expected",
@@ -145,7 +141,7 @@ async def test_llm_planning(llm_agent, mock_cache):
     assert "Plan created successfully" in result["results"][0]["result"]
 
     # Verify plan was stored in cache
-    plan = json.loads(await mock_cache.get("plan"))
+    plan = await mock_cache.get("plan")
 
     print(plan)
     assert isinstance(plan, list)
@@ -194,7 +190,7 @@ async def test_llm_execution(llm_agent, mock_cache):
     execution_history = await mock_cache.get("execution_history")
     assert execution_history is not None
     assert len(execution_history) == 1  # Should have one entry after first execution
-    
+
     history_entry = execution_history[0]
     assert "step" in history_entry
     assert "execution_result" in history_entry
@@ -250,16 +246,23 @@ async def test_missing_api_key():
     """Test that modules properly handle missing API key."""
     if "ANTHROPIC_API_KEY" in os.environ:
         del os.environ["ANTHROPIC_API_KEY"]
-    
-    with pytest.raises(ValueError, match="ANTHROPIC_API_KEY environment variable must be set"):
+
+    with pytest.raises(
+        ValueError, match="ANTHROPIC_API_KEY environment variable must be set"
+    ):
         LLMPlannerModule()
-    
-    with pytest.raises(ValueError, match="ANTHROPIC_API_KEY environment variable must be set"):
+
+    with pytest.raises(
+        ValueError, match="ANTHROPIC_API_KEY environment variable must be set"
+    ):
         LLMExecutorModule()
-    
-    with pytest.raises(ValueError, match="ANTHROPIC_API_KEY environment variable must be set"):
+
+    with pytest.raises(
+        ValueError, match="ANTHROPIC_API_KEY environment variable must be set"
+    ):
         LLMFlowControlModule()
 
+
 @pytest.mark.asyncio
 async def test_error_handling(llm_agent, mock_cache):
     """Test error handling with problematic LLM responses."""
@@ -287,6 +290,8 @@ async def test_error_handling(llm_agent, mock_cache):
 
     # Verify error handling in results
     planner_result = llm_agent.node_results.get("planner", {})
+
+    print(llm_agent.node_results)
     assert planner_result.get("status") in ["success", "error"]
 
     # If planning succeeded despite invalid JSON, executor should handle the error