Merge pull request #2254 from Agenta-AI/litellm-patch

Litellm patch
Agenta-AI · Nov 19, 2024 · a3c559a · a3c559a
2 parents fdc0e82 + d907961
commit a3c559a
Showing 1 changed file with 36 additions and 25 deletions.
diff --git a/agenta-cli/agenta/sdk/litellm/litellm.py b/agenta-cli/agenta/sdk/litellm/litellm.py
@@ -89,16 +89,20 @@ def log_stream_event(
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            result = kwargs.get("complete_streaming_response")
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
 
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+                outputs = {"completion": result}
+                self.span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
 
-            self.span.set_attributes(
-                attributes={"outputs": outputs},
-                namespace="data",
-            )
+            except Exception as e:
+                pass
 
             self.span.set_attributes(
                 attributes={"total": kwargs.get("response_cost")},
@@ -194,16 +198,20 @@ async def async_log_stream_event(
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            result = kwargs.get("complete_streaming_response")
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
 
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+                outputs = {"completion": result}
+                self.span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
 
-            self.span.set_attributes(
-                attributes={"outputs": outputs},
-                namespace="data",
-            )
+            except Exception as e:
+                pass
 
             self.span.set_attributes(
                 attributes={"total": kwargs.get("response_cost")},
@@ -236,17 +244,20 @@ async def async_log_success_event(
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # result = kwargs.get("complete_streaming_response")
-            result = response_obj.choices[0].message.content
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
 
-            outputs = (
-                {"__default__": result} if not isinstance(result, dict) else result
-            )
+                outputs = {"completion": result}
+                self.span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
 
-            self.span.set_attributes(
-                attributes={"outputs": outputs},
-                namespace="data",
-            )
+            except Exception as e:
+                pass
 
             self.span.set_attributes(
                 attributes={"total": kwargs.get("response_cost")},