add typeddict suppoer

Arize-ai · Jan 21, 2025 · 0db4830 · 0db4830
1 parent b399e10
commit 0db4830
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py
@@ -23,6 +23,7 @@
     Type,
     TypeVar,
     Union,
+    _TypedDictMeta,
     cast,
     get_args,
     get_origin,
@@ -1176,6 +1177,14 @@ def _get_jsonschema_type(annotation_type: type) -> Dict[str, Any]:
             _, value_type = annotation_type_args
             result["additionalProperties"] = _get_jsonschema_type(value_type)
         return result
+    is_typed_dict_type = isinstance(annotation_type, _TypedDictMeta)
+    if is_typed_dict_type:
+        result = {"type": "object"}
+        properties = {}
+        for field_name, field_type in annotation_type.__annotations__.items():
+            properties[field_name] = _get_jsonschema_type(field_type)
+        result["properties"] = properties
+        return result
     if (
         pydantic is not None
         and isinstance(annotation_type, type)

diff --git a/python/openinference-instrumentation/tests/test_manual_instrumentation.py b/python/openinference-instrumentation/tests/test_manual_instrumentation.py
@@ -1,7 +1,7 @@
 import json
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, TypedDict, Union
 
 import jsonschema
 import pydantic
@@ -1300,6 +1300,17 @@ class PydanticModel(pydantic.BaseModel):
         bool_param: bool
         any_param: Any
 
+    class TypedDictModel(TypedDict):
+        """
+        typed-dict-description
+        """
+
+        string_param: str
+        int_param: int
+        float_param: float
+        bool_param: bool
+        any_param: Any
+
     AnnotatedWithTypeAlias: TypeAlias = Annotated[str, "This is a description"]
 
     def example_function(  # type: ignore[no-untyped-def]
@@ -1333,6 +1344,7 @@ def example_function(  # type: ignore[no-untyped-def]
         annotated_string_param: Annotated[str, "This is a description"],
         annotated_param_with_type_alias: AnnotatedWithTypeAlias,
         pydantic_model_param: PydanticModel,
+        typed_dict_param: TypedDictModel,
         string_param_with_default: str = "default",
         untyped_param_with_default="default",
     ) -> None:
@@ -1560,6 +1572,24 @@ def example_function(  # type: ignore[no-untyped-def]
                 "title": "PydanticModel",
                 "type": "object",
             },
+            "typed_dict_param": {
+                "type": "object",
+                "properties": {
+                    "string_param": {
+                        "type": "string",
+                    },
+                    "int_param": {
+                        "type": "integer",
+                    },
+                    "float_param": {
+                        "type": "number",
+                    },
+                    "bool_param": {
+                        "type": "boolean",
+                    },
+                    "any_param": {},
+                },
+            },
             "string_param_with_default": {
                 "default": "default",
                 "type": "string",
@@ -1599,6 +1629,7 @@ def example_function(  # type: ignore[no-untyped-def]
             "annotated_string_param",
             "annotated_param_with_type_alias",
             "pydantic_model_param",
+            "typed_dict_param",
         ],
     }
     jsonschema.Draft7Validator.check_schema(