fix: tests for python 3.9

Author: guillaq

workflowai/core/utils/_tools.py

@@ -1,10 +1,8 @@
 import contextlib
-import datetime
 import inspect
-from enum import Enum
-from typing import Any, Callable, NamedTuple, Optional, cast, get_type_hints
+from typing import Any, Callable, NamedTuple, Optional, get_type_hints
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import TypeAdapter
 
 from workflowai.core.utils._schema_generator import JsonSchemaGenerator
 
@@ -17,14 +15,6 @@ class SchemaDeserializer(NamedTuple):
     deserializer: Optional[Callable[[Any], Any]] = None
 
 
-def _serialize_datetime(x: datetime.datetime) -> str:
-    return x.isoformat()
-
-
-def _deserialize_datetime(x: str) -> datetime.datetime:
-    return datetime.datetime.fromisoformat(x)
-
-
 def _get_type_schema(param_type: type):
     """Convert a Python type to its corresponding JSON schema type.
 
@@ -47,31 +37,11 @@ def _get_type_schema(param_type: type):
     if param_type is bool:
         return SchemaDeserializer({"type": "boolean"})
 
-    if param_type is datetime.datetime:
-        return SchemaDeserializer(
-            {"type": "string", "format": "date-time"},
-            serializer=_serialize_datetime,
-            deserializer=_deserialize_datetime,
-        )
-
-    if inspect.isclass(param_type):
-        if issubclass(param_type, BaseModel):
-            return SchemaDeserializer(
-                schema=param_type.model_json_schema(by_alias=True, schema_generator=JsonSchemaGenerator),
-                serializer=lambda x: cast(BaseModel, x).model_dump(mode="json"),  # pyright: ignore [reportUnknownLambdaType]
-                deserializer=param_type.model_validate,
-            )
-
-        if issubclass(param_type, Enum):
-            if not issubclass(param_type, str):
-                raise ValueError(f"Non string enums are not supported: {param_type}")
-            return SchemaDeserializer({"type": "string", "enum": [e.value for e in param_type]})
-
     # Attempting to build a type adapter with pydantic
     with contextlib.suppress(Exception):
         adapter = TypeAdapter[Any](param_type)
         return SchemaDeserializer(
-            schema=adapter.json_schema(),
+            schema=adapter.json_schema(schema_generator=JsonSchemaGenerator),
             deserializer=adapter.validate_python,  # pyright: ignore [reportUnknownLambdaType]
             serializer=lambda x: adapter.dump_python(x, mode="json"),  # pyright: ignore [reportUnknownLambdaType]
         )

workflowai/core/utils/_tools_test.py

@@ -11,6 +11,14 @@
 
 
 class TestGetTypeSchema:
+    class _BasicEnum(str, Enum):
+        A = "a"
+        B = "b"
+
+    class _BasicModel(BaseModel):
+        a: int
+        b: str
+
     @pytest.mark.parametrize(
         ("param_type", "value"),
         [
@@ -22,6 +30,10 @@ class TestGetTypeSchema:
             (ZoneInfo, ZoneInfo("UTC")),
             (list[int], [1, 2, 3]),
             (dict[str, int], {"a": 1, "b": 2}),
+            (_BasicEnum, _BasicEnum.A),
+            (_BasicModel, _BasicModel(a=1, b="test")),
+            (list[_BasicModel], [_BasicModel(a=1, b="test"), _BasicModel(a=2, b="test2")]),
+            (tuple[int, str], (1, "test")),
         ],
     )
     def test_get_type_schema(self, param_type: Any, value: Any):
@@ -50,6 +62,7 @@ def sample_func(
             age: int,
             height: float,
             is_active: bool,
+            date: datetime,
             mode: TestMode = TestMode.FAST,
         ) -> bool:
             """Sample function for testing"""
@@ -77,8 +90,12 @@ def sample_func(
                     "type": "string",
                     "enum": ["fast", "slow"],
                 },
+                "date": {
+                    "type": "string",
+                    "format": "date-time",
+                },
             },
-            "required": ["name", "age", "height", "is_active"],  # 'mode' is not required
+            "required": ["name", "age", "height", "is_active", "date"],  # 'mode' is not required
         }
         assert output_schema.schema == {
             "type": "boolean",