use Model response types

litellm/main.py

@@ -43,12 +43,10 @@ def completion(
     top_k=40, request_timeout=0, # unused var for old version of OpenAI API
   ):
   try:
-    new_response = ModelResponse()
+    model_response = ModelResponse()
-    print(new_response)
     if azure: # this flag is deprecated, remove once notebooks are also updated.
       custom_llm_provider="azure"
     args = locals()
-    model_response = deepcopy(new_response) # deep copy the default response format so we can mutate it and it's thread-safe. 
     # check if user passed in any of the OpenAI optional params
     optional_params = get_optional_params(
       functions=functions, function_call=function_call, 

litellm/tests/test_completion.py

@@ -62,13 +62,19 @@ def test_completion_claude_stream():
 
 def test_completion_cohere():
     try:
-        response = completion(model="command-nightly", messages=messages, max_tokens=500)
+        response = completion(model="command-nightly", messages=messages, max_tokens=100)
         # Add any assertions here to check the response
         print(response)
+        response_str = response['choices'][0]['message']['content']
+        print(f"str response{response_str}")
+        response_str_2 = response.choices[0].message.content
+        if type(response_str) != str:
+            pytest.fail(f"Error occurred: {e}")
+        if type(response_str_2) != str:
+            pytest.fail(f"Error occurred: {e}")
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
-
 def test_completion_cohere_stream():
     try:
         messages = [
@@ -85,6 +91,15 @@ def test_completion_cohere_stream():
 def test_completion_openai():
     try:
         response = completion(model="gpt-3.5-turbo", messages=messages)
+
+        response_str = response['choices'][0]['message']['content']
+        response_str_2 = response.choices[0].message.content
+        print(response_str)
+        print(response_str_2)
+        if type(response_str) != str:
+            pytest.fail(f"Error occurred: {e}")
+        if type(response_str_2) != str:
+            pytest.fail(f"Error occurred: {e}")
         # Add any assertions here to check the response
         print(response)
     except Exception as e:
@@ -197,18 +212,21 @@ def test_completion_replicate_stability_stream():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
-# def test_completion_replicate_stability():
+def test_completion_replicate_stability():
-#     model_name = "stability-ai/stablelm-tuned-alpha-7b:c49dae362cbaecd2ceabb5bd34fdb68413c4ff775111fea065d259d577757beb"
+    model_name = "stability-ai/stablelm-tuned-alpha-7b:c49dae362cbaecd2ceabb5bd34fdb68413c4ff775111fea065d259d577757beb"
-#     try:
+    try:
-#         response = completion(model=model_name, messages=messages, custom_llm_provider="replicate")
+        response = completion(model=model_name, messages=messages, custom_llm_provider="replicate")
-#         # Add any assertions here to check the response
+        # Add any assertions here to check the response
-#         print(response)
+        response_str = response['choices'][0]['message']['content']
-#         for result in response:
+        response_str_2 = response.choices[0].message.content
-#             print(result)
+        print(response_str)
-#         print(response)
+        print(response_str_2)
-#     except Exception as e:
+        if type(response_str) != str:
-#         pytest.fail(f"Error occurred: {e}")
+            pytest.fail(f"Error occurred: {e}")
-# test_completion_replicate_stability()
+        if type(response_str_2) != str:
+            pytest.fail(f"Error occurred: {e}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
 ######## Test TogetherAI ########
 def test_completion_together_ai():

litellm/utils.py

@@ -5,6 +5,7 @@ import litellm, openai
 import random, uuid, requests
 import datetime, time
 import tiktoken
+from pkg_resources import DistributionNotFound, VersionConflict
 encoding = tiktoken.get_encoding("cl100k_base")
 from .integrations.helicone import HeliconeLogger
 from .integrations.aispend import AISpendLogger
@@ -12,6 +13,7 @@ from .integrations.berrispend import BerriSpendLogger
 from .integrations.supabase import Supabase
 from openai.error import OpenAIError as OriginalError
 from .exceptions import AuthenticationError, InvalidRequestError, RateLimitError, ServiceUnavailableError, OpenAIError
+from typing import List, Dict, Union
 ####### ENVIRONMENT VARIABLES ###################
 dotenv.load_dotenv() # Loading env variables using dotenv
 sentry_sdk_instance = None
@@ -48,20 +50,49 @@ local_cache = {}
 #  'usage': {'prompt_tokens': 18, 'completion_tokens': 23, 'total_tokens': 41}
 # }
 
+class Message:
+    def __init__(self):
+        self.content: str = "default"
+        self.role: str = "assistant"
+
+    def __getitem__(self, key):
+      return getattr(self, key)
+
+    def __setitem__(self, key, value):
+      setattr(self, key, value)
+
+    def __iter__(self):
+        return iter(vars(self))
+
+    def __str__(self):
+        result = f"{{\n  'role': '{self.role}',\n  'content': \"{self.content}\"\n}}"
+        return result
+
+class Choices:
+    def __init__(self):
+        self.finish_reason: str = "stop"
+        self.index: int = 0
+        self.message: Message = Message()
+
+    def __getitem__(self, key):
+      return getattr(self, key)
+
+    def __setitem__(self, key, value):
+      setattr(self, key, value)
+
+    def __iter__(self):
+        return iter(vars(self))
+
+    def __str__(self):
+        result = f"{{\n  'finish_reason': '{self.finish_reason}',\n  'index': {self.index},\n  'message': {self.message}\n}}"
+        return result
+
 class ModelResponse:
     def __init__(self):
-        self.choices = [
+        self.choices: List[Choices] = [Choices()]
-          {
+        self.created: str = None
-            "finish_reason": "stop",
+        self.model: str = None
-            "index": 0,
+        self.usage: Dict[str, Union[int, None]] = {
-            "message": {
-                "role": "assistant"
-            }
-          }
-        ]
-        self.created = None
-        self.model = None
-        self.usage = {
             "prompt_tokens": None,
             "completion_tokens": None,
             "total_tokens": None
@@ -77,15 +108,10 @@ class ModelResponse:
         return iter(vars(self))
 
     def __str__(self):
-      result = f"""
+      choices_str = ",\n".join(str(choice) for choice in self.choices)
-                {self.choices}\n
+      result = f"{{\n  'choices': [\n{choices_str}\n  ],\n  'created': {self.created},\n  'model': '{self.model}',\n  'usage': {self.usage}\n}}"
-                {self.created}\n
-                {self.model}\n
-                {self.usage}\n
-                """
       return result
 
-
 def print_verbose(print_statement):
   if litellm.set_verbose:
     print(f"LiteLLM: {print_statement}")