From 13cc5d1c2c49d7c566f6b27f28252ac66ea2128c Mon Sep 17 00:00:00 2001
From: Lou Kratz <219901029+loukratz-bv@users.noreply.github.com>
Date: Wed, 4 Feb 2026 15:46:27 -0500
Subject: [PATCH] fix: #116 Can't Use Domain Types in Tools

---
 .../converters/tools_from_module.py           | 25 +++++++++
 tests/application/test_tools_from_module.py   | 55 +++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/qtype/application/converters/tools_from_module.py b/qtype/application/converters/tools_from_module.py
index 193a940c..4e31d6f7 100644
--- a/qtype/application/converters/tools_from_module.py
+++ b/qtype/application/converters/tools_from_module.py
@@ -1,5 +1,6 @@
 import importlib
 import inspect
+import types
 from typing import Any, Type, Union, get_args, get_origin
 
 from pydantic import BaseModel
@@ -268,6 +269,30 @@ def _map_python_type_to_variable_type(
 
     # Check for generic types like list[str], list[int], etc.
     origin = get_origin(python_type)
+
+    # Handle Union types (including Optional which is Union[T, None])
+    # In Python 3.10+, Type | None creates a types.UnionType
+    is_union = origin is Union or isinstance(python_type, types.UnionType)
+
+    if is_union:
+        args = get_args(python_type)
+        # Filter out None to find the actual type
+        non_none_types = [t for t in args if t is not type(None)]
+
+        if len(non_none_types) == 1:
+            # This is an Optional type (Union[T, None] or T | None)
+            # Recursively map the non-None type
+            return _map_python_type_to_variable_type(
+                non_none_types[0],
+                custom_type_registry,
+                custom_type_models,
+            )
+        else:
+            # Multiple non-None types in union - not currently supported
+            raise ValueError(
+                f"Union types with multiple non-None types are not supported: {python_type}"
+            )
+
     if origin is list:
         # Handle list[T] annotations
         args = get_args(python_type)
diff --git a/tests/application/test_tools_from_module.py b/tests/application/test_tools_from_module.py
index baa31443..71f3e147 100644
--- a/tests/application/test_tools_from_module.py
+++ b/tests/application/test_tools_from_module.py
@@ -21,6 +21,7 @@
     tools_from_module,
 )
 from qtype.base.types import PrimitiveTypeEnum
+from qtype.dsl.domain_types import SearchResult
 from qtype.dsl.model import PythonFunctionTool
 
 
@@ -346,3 +347,57 @@ def test_map_python_type_to_type_str(python_type, expected):
     """Test type to string mapping."""
     result = _map_python_type_to_type_str(python_type, {}, {})
     assert result == expected
+
+
+def test_tools_from_module_with_optional_domain_type(temp_module):
+    """Test function with optional domain type parameter (Union[Type, None])."""
+    module_name = temp_module(
+        """
+        from qtype.dsl.domain_types import SearchResult
+
+        def process_result(result: SearchResult | None = None) -> str:
+            '''Process a search result.'''
+            if result is None:
+                return "No result"
+            return str(result.doc_id)
+        """,
+        "optional_domain_module",
+    )
+
+    tools, custom_types = tools_from_module(module_name)
+
+    assert len(tools) == 1
+    tool = tools[0]
+    assert tool.name == "process_result"
+    assert tool.inputs is not None
+    assert len(tool.inputs) == 1
+
+    # Check that the parameter is properly typed as optional SearchResult
+    result_param = next((v for v in tool.inputs if v.id == "result"), None)
+    assert result_param is not None
+    assert result_param.type == SearchResult
+    assert result_param.optional
+
+
+def test_map_python_type_to_variable_type_union_with_none():
+    """Test mapping Union type with None to optional variable type."""
+    from typing import Union
+
+    # Test with domain type
+    result = _map_python_type_to_variable_type(
+        Union[SearchResult, None], {}, {}
+    )
+    assert result == "SearchResult"
+
+    # Test with primitive type
+    result = _map_python_type_to_variable_type(Union[str, None], {}, {})
+    assert result == PrimitiveTypeEnum.text
+
+    # Test with Pydantic model
+    custom_type_registry = {}
+    custom_type_models = {}
+    result = _map_python_type_to_variable_type(
+        Union[SampleModel, None], custom_type_registry, custom_type_models
+    )
+    assert result == "SampleModel"
+    assert "SampleModel" in custom_type_registry