openai mock correcting

Author: virrius

sgr_agent_core/agents/sgr_tool_calling_agent.py

@@ -66,7 +66,7 @@ async def _reasoning_phase(self) -> ReasoningTool:
                 ],
             }
         )
-        tool_call_result = await reasoning(self._context)
+        tool_call_result = await reasoning(self._context, self.config)
         self.streaming_generator.add_tool_call(
             f"{self._context.iteration}-reasoning", reasoning.tool_name, tool_call_result
         )

tests/test_agent_e2e.py

@@ -16,29 +16,69 @@
 
 
 class MockStream:
-    """Mock OpenAI stream object."""
+    """Mock OpenAI stream object that emulates OpenAI streaming API.
+
+    This mock properly handles:
+    - Context manager protocol (async with)
+    - Stream iteration (async for event in stream)
+    - Final completion retrieval with parsed_arguments support
+    """
 
     def __init__(self, final_completion_data: dict):
+        """Initialize mock stream with final completion data.
+
+        Args:
+            final_completion_data: Dictionary containing:
+                - content: Optional message content
+                - tool_calls: List of tool call objects (already with parsed_arguments set)
+        """
         self._final_completion_data = final_completion_data
+        self._iterated = False
 
     async def __aenter__(self):
+        """Enter context manager."""
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager."""
         pass
 
     def __aiter__(self):
+        """Return iterator for stream events."""
         return self
 
     async def __anext__(self):
+        """Return next stream event (empty iterator for simplicity).
+
+        In real OpenAI API, this would yield chunk events. For testing,
+        we return empty iterator since the code handles missing chunks
+        gracefully.
+        """
+        if self._iterated:
+            raise StopAsyncIteration
+        self._iterated = True
         raise StopAsyncIteration
 
     async def get_final_completion(self) -> ChatCompletion:
+        """Get final completion with parsed tool call arguments or structured
+        output.
+
+        Supports both formats:
+        - Structured output: message.parsed (for SGRAgent)
+        - Function calling: tool_calls[0].function.parsed_arguments (for SGRToolCallingAgent)
+
+        Returns:
+            ChatCompletion object with appropriate parsed data
+        """
+        tool_calls = self._final_completion_data.get("tool_calls", [])
+
         message = ChatCompletionMessage(
             role="assistant",
             content=self._final_completion_data.get("content"),
-            tool_calls=self._final_completion_data.get("tool_calls"),
+            tool_calls=tool_calls if tool_calls else None,
         )
+
+        # Support structured output format (SGRAgent uses message.parsed)
         if "parsed" in self._final_completion_data:
             setattr(message, "parsed", self._final_completion_data["parsed"])
 
@@ -155,6 +195,15 @@ def mock_stream(**kwargs):
 
 
 def create_mock_openai_client_for_sgr_tool_calling_agent(action_tool_1: Type, action_tool_2: Type) -> AsyncOpenAI:
+    """Create a mock OpenAI client for SGRToolCallingAgent tests.
+
+    Args:
+        action_tool_1: First action tool to return (e.g., AdaptPlanTool)
+        action_tool_2: Second action tool to return (e.g., FinalAnswerTool)
+
+    Returns:
+        Mocked AsyncOpenAI client configured for SGRToolCallingAgent execution cycle
+    """
     client = Mock(spec=AsyncOpenAI)
 
     reasoning_tools = [
@@ -196,32 +245,41 @@ def create_mock_openai_client_for_sgr_tool_calling_agent(action_tool_1: Type, ac
     action_count = {"count": 0}
 
     def mock_stream(**kwargs):
-        is_reasoning = (
-            "tool_choice" in kwargs
-            and isinstance(kwargs.get("tool_choice"), dict)
-            and kwargs["tool_choice"].get("function", {}).get("name") == ReasoningTool.tool_name
-        )
+        """Mock stream function that determines reasoning vs action phase.
+
+        Checks the first tool name in the tools list to determine the
+        phase.
+        """
+        tools_param = kwargs.get("tools", [])
+
+        # Check first tool name - if it's ReasoningTool, it's reasoning phase
+        tool_name = None
+        if isinstance(tools_param, list) and tools_param:
+            first_tool = tools_param[0]
+            if isinstance(first_tool, dict):
+                tool_name = first_tool.get("function", {}).get("name")
+
+        is_reasoning = tool_name == ReasoningTool.tool_name
 
         if is_reasoning:
             reasoning_count["count"] += 1
             tool = reasoning_tools[reasoning_count["count"] - 1]
-            call_id = f"reasoning_{reasoning_count['count']}"
+            call_id = f"{reasoning_count['count']}-reasoning"
         else:
-            tools_param = kwargs.get("tools")
-            if tools_param is not None and not isinstance(tools_param, list):
+            # Validate that tools is a list (for action phase)
+            if not isinstance(tools_param, list):
                 raise TypeError(
                     f"SGRToolCallingAgent._prepare_tools() must return a list, "
                     f"but got {type(tools_param).__name__}. "
                     f"Override _prepare_tools() to return list instead of NextStepToolStub."
                 )
             action_count["count"] += 1
             tool = action_tools[action_count["count"] - 1]
-            call_id = f"action_{action_count['count']}"
+            call_id = f"{action_count['count']}-action"
 
         return MockStream(
             final_completion_data={
                 "content": None,
-                "role": "assistant",
                 "tool_calls": [_create_tool_call(tool, call_id)],
             }
         )