From 9c3a1124d7bdcd448871e0fb6cdc8bcbc4086394 Mon Sep 17 00:00:00 2001
From: Vivek Aditya <vivek@athina.ai>
Date: Wed, 17 Apr 2024 00:06:16 +0530
Subject: [PATCH] Athina callback now handles streaming mode

---
 litellm/integrations/athina.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/litellm/integrations/athina.py b/litellm/integrations/athina.py
index 897cf6c8d53..97d84c57433 100644
--- a/litellm/integrations/athina.py
+++ b/litellm/integrations/athina.py
@@ -1,6 +1,5 @@
 import datetime
 
-
 class AthinaLogger:
     def __init__(self):
         import os
@@ -17,7 +16,18 @@ class AthinaLogger:
         import json
         import traceback
         try:
-            response_json = response_obj.model_dump() if response_obj else {}
+            is_stream = kwargs.get("stream", False)
+            if is_stream:
+                if "complete_streaming_response" in kwargs:
+                    # Log the completion response in streaming mode
+                    completion_response = kwargs["complete_streaming_response"]
+                    response_json = completion_response.model_dump() if completion_response else {}
+                else:
+                    # Skip logging if the completion response is not available
+                    return
+            else:
+                # Log the completion response in non streaming mode
+                response_json = response_obj.model_dump() if response_obj else {}
             data = {
                 "language_model_id": kwargs.get("model"),
                 "request": kwargs,