reasoning tool as parameter (#180)

Author: virrius

sgr_agent_core/agents/iron_agent.py

@@ -11,7 +11,7 @@
 from sgr_agent_core.next_step_tool import NextStepToolsBuilder
 from sgr_agent_core.services.registry import ToolRegistry
 from sgr_agent_core.services.tool_instantiator import ToolInstantiator
-from sgr_agent_core.tools import BaseTool, ReasoningTool, ToolNameSelectorStub
+from sgr_agent_core.tools import BaseTool, ReasoningTool, SystemBaseTool, ToolNameSelectorStub
 
 
 class IronAgent(BaseAgent):
@@ -34,7 +34,9 @@ def __init__(
         openai_client: AsyncOpenAI,
         agent_config: AgentConfig,
         toolkit: list[Type[BaseTool]],
+        *,
         def_name: str | None = None,
+        reasoning_tool_cls: type[SystemBaseTool] = ReasoningTool,
         **kwargs: dict,
     ):
         super().__init__(
@@ -45,6 +47,7 @@ def __init__(
             def_name=def_name,
             **kwargs,
         )
+        self.ReasoningTool: type[SystemBaseTool] = reasoning_tool_cls
 
     def _log_tool_instantiator(
         self,
@@ -144,7 +147,7 @@ async def _prepare_tools(self) -> Type[ToolNameSelectorStub]:
         """Prepare available tools for the current agent state and progress."""
         if self._context.iteration >= self.config.execution.max_iterations:
             raise RuntimeError("Max iterations reached")
-        return NextStepToolsBuilder.build_NextStepToolSelector(self.toolkit)
+        return NextStepToolsBuilder.build_NextStepToolSelector(self.toolkit, base_reasoning_cls=self.ReasoningTool)
 
     async def _reasoning_phase(self) -> ReasoningTool:
         """Call LLM to get ReasoningTool with selected tool name."""
@@ -153,8 +156,8 @@ async def _reasoning_phase(self) -> ReasoningTool:
         tool_selector_model = await self._prepare_tools()
         reasoning = await self._generate_tool(tool_selector_model, messages)
 
-        if not isinstance(reasoning, ReasoningTool):
-            raise ValueError("Expected ReasoningTool instance")
+        if not isinstance(reasoning, self.ReasoningTool):
+            raise ValueError(f"Expected {self.ReasoningTool.__name__} instance")
 
         # Log reasoning
         self._log_reasoning(reasoning)

sgr_agent_core/agents/sgr_agent.py

@@ -8,6 +8,8 @@
 from sgr_agent_core.tools import (
     BaseTool,
     NextStepToolStub,
+    ReasoningTool,
+    SystemBaseTool,
 )
 
 
@@ -22,7 +24,9 @@ def __init__(
         openai_client: AsyncOpenAI,
         agent_config: AgentConfig,
         toolkit: list[Type[BaseTool]],
+        *,
         def_name: str | None = None,
+        reasoning_tool_cls: type[SystemBaseTool] = ReasoningTool,
         **kwargs: dict,
     ):
         super().__init__(
@@ -33,11 +37,12 @@ def __init__(
             def_name=def_name,
             **kwargs,
         )
+        self.ReasoningTool: type[SystemBaseTool] = reasoning_tool_cls
 
     async def _prepare_tools(self) -> Type[NextStepToolStub]:
         """Prepare available tools for the current agent state and progress."""
         tools = set(self.toolkit)
-        return NextStepToolsBuilder.build_NextStepTools(list(tools))
+        return NextStepToolsBuilder.build_NextStepTools(list(tools), base_reasoning_cls=self.ReasoningTool)
 
     async def _reasoning_phase(self) -> NextStepToolStub:
         phase_id = f"{self._context.iteration}-reasoning"

sgr_agent_core/agents/sgr_tool_calling_agent.py

@@ -9,6 +9,7 @@
     BaseTool,
     FinalAnswerTool,
     ReasoningTool,
+    SystemBaseTool,
 )
 
 
@@ -24,7 +25,9 @@ def __init__(
         openai_client: AsyncOpenAI,
         agent_config: AgentConfig,
         toolkit: list[Type[BaseTool]],
+        *,
         def_name: str | None = None,
+        reasoning_tool_cls: type[SystemBaseTool] = ReasoningTool,
         **kwargs: dict,
     ):
         super().__init__(
@@ -36,12 +39,13 @@ def __init__(
             **kwargs,
         )
         self.tool_choice: Literal["required"] = "required"
+        self.ReasoningTool: type[SystemBaseTool] = reasoning_tool_cls
 
     async def _reasoning_phase(self) -> ReasoningTool:
         phase_id = f"{self._context.iteration}-reasoning"
         async with self.openai_client.chat.completions.stream(
             messages=await self._prepare_context(),
-            tools=[pydantic_function_tool(ReasoningTool, name=ReasoningTool.tool_name)],
+            tools=[pydantic_function_tool(self.ReasoningTool, name=self.ReasoningTool.tool_name)],
             tool_choice=self.tool_choice,
             **self.config.llm.to_openai_client_kwargs(),
         ) as stream:

sgr_agent_core/base_tool.py

@@ -2,7 +2,7 @@
 
 import json
 import logging
-from typing import TYPE_CHECKING, ClassVar, Self
+from typing import TYPE_CHECKING, ClassVar, Self, TypeVar
 
 from fastmcp import Client
 from pydantic import BaseModel
@@ -54,6 +54,9 @@ class SystemBaseTool(BaseTool):
     isSystemTool: ClassVar[bool] = True
 
 
+ReasoningToolStubType = TypeVar("ReasoningToolStubType", bound=SystemBaseTool)
+
+
 class MCPBaseTool(BaseTool):
     """Base model for MCP Tool schema."""
 

sgr_agent_core/next_step_tool.py

@@ -8,27 +8,32 @@
 
 from pydantic import BaseModel, Field, create_model
 
-from sgr_agent_core.base_tool import BaseTool
+from sgr_agent_core.base_tool import BaseTool, SystemBaseTool
 from sgr_agent_core.tools.reasoning_tool import ReasoningTool
 
 logger = logging.getLogger(__name__)
 
 T = TypeVar("T", bound=BaseTool)
 
 
-class NextStepToolStub(ReasoningTool, ABC):
-    """SGR Core - Determines the next reasoning step with adaptive planning, choosing appropriate tool
-    (!) Stub class for correct autocomplete. Use NextStepToolsBuilder"""
+class NextStepToolStub(SystemBaseTool, ABC):
+    """SGR Core - Determines the next reasoning step with adaptive planning, choosing appropriate tool.
+
+    (!) Stub class for correct autocomplete. Use NextStepToolsBuilder.
+    The actual base reasoning class is injected at build time.
+    """
 
     function: T = Field(description="Select the appropriate tool for the next step")
 
 
-class ToolNameSelectorStub(ReasoningTool, ABC):
-    """Stub class for tool name selection that inherits from ReasoningTool.
+class ToolNameSelectorStub(SystemBaseTool, ABC):
+    """Stub class for tool name selection.
 
     Used by IronAgent to select tool name as part of reasoning phase.
     (!) Stub class for correct autocomplete. Use
-    NextStepToolsBuilder.build_NextStepToolSelector
+    NextStepToolsBuilder.build_NextStepToolSelector with
+    base_reasoning_cls. The actual base reasoning class is injected at
+    build time.
     """
 
     function_name_choice: str = Field(description="Select the name of the tool to use")
@@ -70,21 +75,43 @@ def _create_tool_types_union(cls, tools_list: list[Type[T]]) -> Type:
         return Annotated[union, Field()]
 
     @classmethod
-    def build_NextStepTools(cls, tools_list: list[Type[T]]) -> Type[NextStepToolStub]:  # noqa
-        """Build a model with all NextStepTool args."""
+    def build_NextStepTools(  # noqa
+        cls,
+        tools_list: list[Type[T]],
+        base_reasoning_cls: type[ReasoningTool] = ReasoningTool,
+    ) -> Type[NextStepToolStub]:
+        """Build a model with all NextStepTool args.
+
+        Args:
+            tools_list: List of tool classes to include in the union.
+            base_reasoning_cls: Pydantic model class used as the base for the
+                reasoning schema sent to the LLM via Structured Output. Defaults
+                to ReasoningTool. Pass a subclass to extend or override the
+                reasoning schema.
+        """
         return create_model(
             "NextStepTools",
-            __base__=NextStepToolStub,
+            __base__=base_reasoning_cls,
             function=(
                 cls._create_tool_types_union(tools_list),
                 Field(description="Select and fill parameters of the appropriate tool for the next step"),
             ),
         )
 
     @classmethod
-    def build_NextStepToolSelector(cls, tools_list: list[Type[T]]) -> Type[ToolNameSelectorStub]:
-        """Build a model for selecting tool name."""
-        # Extract tool names and descriptions
+    def build_NextStepToolSelector(  # noqa
+        cls,
+        tools_list: list[Type[T]],
+        base_reasoning_cls: type[SystemBaseTool] = ReasoningTool,
+    ) -> Type[ToolNameSelectorStub]:
+        """Build a model for selecting tool name.
+
+        Args:
+            tools_list: List of tool classes whose names form the allowed choices.
+            base_reasoning_cls: Pydantic model class used as the base for the
+                reasoning schema. Defaults to ReasoningTool. Pass a subclass to
+                extend or override the reasoning schema.
+        """
         tool_names = [tool.tool_name for tool in tools_list]
 
         if len(tool_names) == 1:
@@ -98,7 +125,7 @@ def build_NextStepToolSelector(cls, tools_list: list[Type[T]]) -> Type[ToolNameS
         # Create model dynamically, inheriting from ToolNameSelectorStub (which inherits from ReasoningTool)
         model_class = create_model(
             "NextStepToolSelector",
-            __base__=ToolNameSelectorStub,
+            __base__=base_reasoning_cls,
             function_name_choice=(literal_type, Field(description="Choose the name for the best tool to use")),
         )
         model_class.tool_name = "nextsteptoolselector"  # type: ignore

sgr_agent_core/tools/__init__.py

@@ -1,4 +1,4 @@
-from sgr_agent_core.base_tool import BaseTool, MCPBaseTool, SystemBaseTool
+from sgr_agent_core.base_tool import BaseTool, MCPBaseTool, ReasoningToolStubType, SystemBaseTool
 from sgr_agent_core.next_step_tool import (
     NextStepToolsBuilder,
     NextStepToolStub,
@@ -19,6 +19,7 @@
     "BaseTool",
     "MCPBaseTool",
     "SystemBaseTool",
+    "ReasoningToolStubType",
     "NextStepToolStub",
     "ToolNameSelectorStub",
     "NextStepToolsBuilder",

tests/test_agent_e2e.py

@@ -364,3 +364,129 @@ async def test_sgr_tool_calling_agent_full_execution_cycle():
 
     assert result is not None
     _assert_agent_completed(agent)
+
+
+@pytest.mark.asyncio
+async def test_sgr_tool_calling_agent_custom_reasoning_tool_is_used():
+    """Custom ReasoningTool is actually passed to OpenAI in _reasoning_phase.
+
+    Verifies that self.ReasoningTool is forwarded to
+    pydantic_function_tool() instead of the hardcoded base
+    ReasoningTool.
+    """
+    from pydantic import Field as PydanticField
+
+    class CustomReasoningTool(ReasoningTool):
+        confidence: float = PydanticField(default=0.5, description="Confidence in the decision")
+
+    captured_reasoning_tool_names: list[str] = []
+
+    reasoning_instance = CustomReasoningTool(
+        reasoning_steps=["Analyze", "Decide"],
+        current_situation="Test situation",
+        plan_status="On track",
+        enough_data=False,
+        remaining_steps=["Finalize"],
+        task_completed=False,
+        confidence=0.9,
+    )
+    final_answer_instance = FinalAnswerTool(
+        reasoning="Done",
+        completed_steps=["Step 1"],
+        answer="Final answer to the research task",
+        status=AgentStatesEnum.COMPLETED,
+    )
+
+    client = Mock(spec=AsyncOpenAI)
+
+    def mock_stream(**kwargs):
+        tools_param = kwargs.get("tools", [])
+        tool_name = None
+        if tools_param and isinstance(tools_param, list) and isinstance(tools_param[0], dict):
+            tool_name = tools_param[0].get("function", {}).get("name")
+
+        if tool_name == CustomReasoningTool.tool_name:
+            captured_reasoning_tool_names.append(tool_name)
+            return MockStream({"content": None, "tool_calls": [_create_tool_call(reasoning_instance, "call-r")]})
+
+        return MockStream({"content": None, "tool_calls": [_create_tool_call(final_answer_instance, "call-a")]})
+
+    client.chat.completions.stream = Mock(side_effect=mock_stream)
+
+    agent = SGRToolCallingAgent(
+        task_messages=[{"role": "user", "content": "Test task"}],
+        openai_client=client,
+        agent_config=_create_test_agent_config(),
+        toolkit=[FinalAnswerTool],
+        reasoning_tool_cls=CustomReasoningTool,
+    )
+
+    result = await agent.execute()
+
+    assert result is not None
+    assert agent._context.state == AgentStatesEnum.COMPLETED
+    assert len(captured_reasoning_tool_names) >= 1, "Custom ReasoningTool was never passed to OpenAI"
+    assert captured_reasoning_tool_names[0] == CustomReasoningTool.tool_name
+
+
+@pytest.mark.asyncio
+async def test_sgr_agent_custom_reasoning_tool_is_used():
+    """Custom ReasoningTool is used as SO base in SGRAgent._reasoning_phase.
+
+    Verifies that response_format passed to OpenAI is built on top of
+    the custom ReasoningTool subclass rather than the default one.
+    """
+    from pydantic import Field as PydanticField
+
+    class CustomReasoningTool(ReasoningTool):
+        confidence: float = PydanticField(default=0.5, description="Confidence in the decision")
+
+    captured_response_formats: list[type] = []
+
+    client = Mock(spec=AsyncOpenAI)
+
+    def mock_stream(**kwargs):
+        response_format = kwargs.get("response_format")
+        if response_format is not None:
+            captured_response_formats.append(response_format)
+
+        NextStepTools = NextStepToolsBuilder.build_NextStepTools(
+            [FinalAnswerTool],
+            base_reasoning_cls=CustomReasoningTool,
+        )
+        response = NextStepTools(
+            reasoning_steps=["Step 1", "Step 2"],
+            current_situation="Test",
+            plan_status="Ok",
+            enough_data=True,
+            remaining_steps=["Finalize"],
+            task_completed=True,
+            confidence=0.8,
+            function={
+                "tool_name_discriminator": FinalAnswerTool.tool_name,
+                "reasoning": "Done",
+                "completed_steps": ["Step 1"],
+                "answer": "Final answer to the research task",
+                "status": AgentStatesEnum.COMPLETED,
+            },
+        )
+        return MockStream({"parsed": response})
+
+    client.chat.completions.stream = Mock(side_effect=mock_stream)
+
+    agent = SGRAgent(
+        task_messages=[{"role": "user", "content": "Test task"}],
+        openai_client=client,
+        agent_config=_create_test_agent_config(),
+        toolkit=[FinalAnswerTool],
+        reasoning_tool_cls=CustomReasoningTool,
+    )
+
+    result = await agent.execute()
+
+    assert result is not None
+    assert agent._context.state == AgentStatesEnum.COMPLETED
+    assert len(captured_response_formats) >= 1, "response_format was never passed to OpenAI"
+    assert issubclass(
+        captured_response_formats[0], CustomReasoningTool
+    ), f"response_format {captured_response_formats[0]} is not a subclass of CustomReasoningTool"