From 9b46412279b887a60b8e82ca4db0d7da81072f7f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 8 Jan 2024 12:05:20 +0530
Subject: [PATCH] fix(utils.py): fix logging for text completion streaming

---
 litellm/tests/test_custom_callback_input.py | 39 +++++++++++++++++++++
 litellm/utils.py                            |  4 ++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py
index 8f28e86d9ba..0fb69b64513 100644
--- a/litellm/tests/test_custom_callback_input.py
+++ b/litellm/tests/test_custom_callback_input.py
@@ -545,6 +545,45 @@ async def test_async_chat_bedrock_stream():
 
 # asyncio.run(test_async_chat_bedrock_stream())
 
+# Text Completion 
+        
+## Test OpenAI text completion + Async
+@pytest.mark.asyncio
+async def test_async_text_completion_openai_stream():
+    try:
+        customHandler = CompletionCustomHandler()
+        litellm.callbacks = [customHandler]
+        response = await litellm.atext_completion(
+            model="gpt-3.5-turbo",
+            prompt="Hi 👋 - i'm async text completion openai",
+        )
+        # test streaming
+        response = await litellm.atext_completion(
+            model="gpt-3.5-turbo",
+            prompt="Hi 👋 - i'm async text completion openai",
+            stream=True,
+        )
+        async for chunk in response:
+            print(f"chunk: {chunk}")
+            continue
+        ## test failure callback
+        try:
+            response = await litellm.atext_completion(
+                model="gpt-3.5-turbo",
+                prompt="Hi 👋 - i'm async text completion openai",
+                stream=True,
+                api_key="my-bad-key",
+            )
+            async for chunk in response:
+                continue
+        except:
+            pass
+        time.sleep(1)
+        print(f"customHandler.errors: {customHandler.errors}")
+        assert len(customHandler.errors) == 0
+        litellm.callbacks = []
+    except Exception as e:
+        pytest.fail(f"An exception occurred: {str(e)}")
 
 # EMBEDDING
 ## Test OpenAI + Async
diff --git a/litellm/utils.py b/litellm/utils.py
index 0a8c410ea04..da1dba897cc 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -710,7 +710,7 @@ class CallTypes(Enum):
     aimage_generation = "aimage_generation"
 
 
-# Logging function -> log the exact model details + what's being sent | Non-Blocking
+# Logging function -> log the exact model details + what's being sent | Non-BlockingP
 class Logging:
     global supabaseClient, liteDebuggerClient, promptLayerLogger, weightsBiasesLogger, langsmithLogger, capture_exception, add_breadcrumb, llmonitorLogger
 
@@ -729,6 +729,8 @@ class Logging:
             raise ValueError(
                 f"Invalid call_type {call_type}. Allowed values: {allowed_values}"
             )
+        if messages is not None and isinstance(messages, str):
+            messages = [{"role": "user", "content": messages}] # convert text completion input to the chat completion format
         self.model = model
         self.messages = messages
         self.stream = stream