From 9d5148fd8e08a9415debe34e2069de963867e6ad Mon Sep 17 00:00:00 2001
From: Nate Mar <nate@arize.com>
Date: Mon, 13 Jan 2025 15:29:21 -0800
Subject: [PATCH] Update __init__.py

---
 .../src/openinference/instrumentation/litellm/__init__.py      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py b/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py
index 180c695ba..19ef943de 100644
--- a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py
+++ b/python/instrumentation/openinference-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