Merge pull request #172 from BrainBlend-AI/feature/161-input-output-schema-have-unexpected-values

Fix schema properties returning BaseIOSchema instead of custom types

  Add __init_subclass__ hook to capture generic type parameters during class
  definition, fixing issue where input_schema and output_schema properties
  returned BaseIOSchema instead of the specified custom schema types.

  - Implement type parameter capture for inheritance pattern
  - Maintain fallback for dynamic instantiation and edge cases
  - Add focused test coverage for schema resolution
  - Resolves GitHub issue #161

Author: KennyVaneetvelde

atomic-agents/atomic_agents/base/base_tool.py

@@ -1,4 +1,4 @@
-from typing import Optional, Type, get_args
+from typing import Optional, Type, get_args, get_origin
 from abc import ABC, abstractmethod
 from pydantic import BaseModel
 
@@ -47,6 +47,23 @@ def __init__(self, config: BaseToolConfig = BaseToolConfig()):
         """
         self.config = config
 
+    def __init_subclass__(cls, **kwargs):
+        """
+        Hook called when a class is subclassed.
+
+        Captures generic type parameters during class creation and stores them as class attributes
+        to work around the unreliable __orig_class__ attribute in modern Python generic syntax.
+        """
+        super().__init_subclass__(**kwargs)
+        if hasattr(cls, "__orig_bases__"):
+            for base in cls.__orig_bases__:
+                if get_origin(base) is BaseTool:
+                    args = get_args(base)
+                    if len(args) == 2:
+                        cls._input_schema_cls = args[0]
+                        cls._output_schema_cls = args[1]
+                        break
+
     @property
     def input_schema(self) -> Type[InputSchema]:
         """
@@ -55,12 +72,17 @@ def input_schema(self) -> Type[InputSchema]:
         Returns:
             Type[InputSchema]: The input schema class.
         """
+        # Inheritance pattern: MyTool(BaseTool[Schema1, Schema2])
+        if hasattr(self.__class__, "_input_schema_cls"):
+            return self.__class__._input_schema_cls
+
+        # Dynamic instantiation: MockTool[Schema1, Schema2]()
         if hasattr(self, "__orig_class__"):
             TI, _ = get_args(self.__orig_class__)
-        else:
-            TI = BaseIOSchema
+            return TI
 
-        return TI
+        # No type info available: MockTool()
+        return BaseIOSchema
 
     @property
     def output_schema(self) -> Type[OutputSchema]:
@@ -70,12 +92,17 @@ def output_schema(self) -> Type[OutputSchema]:
         Returns:
             Type[OutputSchema]: The output schema class.
         """
+        # Inheritance pattern: MyTool(BaseTool[Schema1, Schema2])
+        if hasattr(self.__class__, "_output_schema_cls"):
+            return self.__class__._output_schema_cls
+
+        # Dynamic instantiation: MockTool[Schema1, Schema2]()
         if hasattr(self, "__orig_class__"):
             _, TO = get_args(self.__orig_class__)
-        else:
-            TO = BaseIOSchema
+            return TO
 
-        return TO
+        # No type info available: MockTool()
+        return BaseIOSchema
 
     @property
     def tool_name(self) -> str:

atomic-agents/tests/base/test_base_tool.py

@@ -94,3 +94,30 @@ def test_base_tool_config_optional_fields():
     config = BaseToolConfig()
     assert hasattr(config, "title")
     assert hasattr(config, "description")
+
+
+# Test for GitHub issue #161 fix: proper schema resolution
+def test_base_tool_schema_resolution():
+    """Test that input_schema and output_schema return correct types (not BaseIOSchema)"""
+
+    class CustomInput(BaseIOSchema):
+        """Custom input schema for testing"""
+
+        name: str
+
+    class CustomOutput(BaseIOSchema):
+        """Custom output schema for testing"""
+
+        result: str
+
+    class TestTool(BaseTool[CustomInput, CustomOutput]):
+        def run(self, params: CustomInput) -> CustomOutput:
+            return CustomOutput(result=f"processed_{params.name}")
+
+    tool = TestTool()
+
+    # These should return the specific types, not BaseIOSchema
+    assert tool.input_schema == CustomInput
+    assert tool.output_schema == CustomOutput
+    assert tool.input_schema != BaseIOSchema
+    assert tool.output_schema != BaseIOSchema