From ac6fb0cbef0d62679d4cdf21e8ea38bf255fc235 Mon Sep 17 00:00:00 2001
From: Willy Douhard <willy.douhard@gmail.com>
Date: Tue, 8 Oct 2024 17:33:32 +0200
Subject: [PATCH] Fix: Literal AI llm completion logging (#6096)

* fix: log llm output

* chore: rename var
---
 litellm/integrations/literal_ai.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/litellm/integrations/literal_ai.py b/litellm/integrations/literal_ai.py
index 0a7f2f675..f83b714e5 100644
--- a/litellm/integrations/literal_ai.py
+++ b/litellm/integrations/literal_ai.py
@@ -2,7 +2,6 @@
 # This file contains the LiteralAILogger class which is used to log steps to the LiteralAI observability platform.
 import asyncio
 import os
-import traceback
 import uuid
 from typing import Optional
 
@@ -10,7 +9,6 @@ import httpx
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
-from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
 from litellm.llms.custom_httpx.http_handler import (
     HTTPHandler,
     get_async_httpx_client,
@@ -92,7 +90,6 @@ class LiteralAILogger(CustomBatchLogger):
         url = f"{self.literalai_api_url}/api/graphql"
         query = self._steps_query_builder(self.log_queue)
         variables = self._steps_variables_builder(self.log_queue)
-
         try:
             response = self.sync_http_handler.post(
                 url=url,
@@ -102,7 +99,6 @@ class LiteralAILogger(CustomBatchLogger):
                 },
                 headers=self.headers,
             )
-            response.raise_for_status()
 
             if response.status_code >= 300:
                 verbose_logger.error(
@@ -170,8 +166,6 @@ class LiteralAILogger(CustomBatchLogger):
                 },
                 headers=self.headers,
             )
-            response.raise_for_status()
-
             if response.status_code >= 300:
                 verbose_logger.error(
                     f"Literal AI Error: {response.status_code} - {response.text}"
@@ -200,6 +194,8 @@ class LiteralAILogger(CustomBatchLogger):
         settings = logging_payload["model_parameters"]
 
         messages = logging_payload["messages"]
+        choices = logging_payload.get("response", {}).get("choices", [])
+        message_completion = choices[0]["message"] if choices else None
         prompt_id = None
         variables = None
 
@@ -238,6 +234,7 @@ class LiteralAILogger(CustomBatchLogger):
                 "duration": (end_time - start_time).total_seconds(),
                 "settings": settings,
                 "messages": messages,
+                "messageCompletion": message_completion,
                 "tools": tools,
             },
         }