Update __init__.py

Arize-ai · Jan 13, 2025 · 9d5148f · 9d5148f
1 parent 0d9ba2f
commit 9d5148f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/...eninference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py b/...eninference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py
@@ -11,6 +11,7 @@
     Mapping,
     Tuple,
     TypeVar,
+    Union,
 )
 
 from openai.types.image import Image
@@ -59,7 +60,7 @@ def is_iterable_of(lst: Iterable[object], tp: T) -> bool:
 
 
 def _get_attributes_from_message_param(
-    message: Mapping[str, Any],
+    message: Union[Mapping[str, Any], litellm.Message],
 ) -> Iterator[Tuple[str, AttributeValue]]:
     if not hasattr(message, "get"):
         return