diff --git a/litellm/main.py b/litellm/main.py index f3fe0f88a..1dd6fc954 100644 --- a/litellm/main.py +++ b/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() - print(new_response) + model_response = ModelResponse() 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, diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index ec6baec1c..79e18a0f3 100644 --- a/litellm/tests/test_completion.py +++ b/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(): -# model_name = "stability-ai/stablelm-tuned-alpha-7b:c49dae362cbaecd2ceabb5bd34fdb68413c4ff775111fea065d259d577757beb" -# try: -# response = completion(model=model_name, messages=messages, custom_llm_provider="replicate") -# # Add any assertions here to check the response -# print(response) -# for result in response: -# print(result) -# print(response) -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") -# test_completion_replicate_stability() +def test_completion_replicate_stability(): + model_name = "stability-ai/stablelm-tuned-alpha-7b:c49dae362cbaecd2ceabb5bd34fdb68413c4ff775111fea065d259d577757beb" + try: + response = completion(model=model_name, messages=messages, custom_llm_provider="replicate") + # Add any assertions here to check the response + 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}") + except Exception as e: + pytest.fail(f"Error occurred: {e}") ######## Test TogetherAI ######## def test_completion_together_ai(): diff --git a/litellm/utils.py b/litellm/utils.py index b27938e90..02c0c49c4 100644 --- a/litellm/utils.py +++ b/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 = [ - { - "finish_reason": "stop", - "index": 0, - "message": { - "role": "assistant" - } - } - ] - self.created = None - self.model = None - self.usage = { + self.choices: List[Choices] = [Choices()] + self.created: str = None + self.model: str = None + self.usage: Dict[str, Union[int, None]] = { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None @@ -77,15 +108,10 @@ class ModelResponse: return iter(vars(self)) def __str__(self): - result = f""" - {self.choices}\n - {self.created}\n - {self.model}\n - {self.usage}\n - """ + choices_str = ",\n".join(str(choice) for choice in self.choices) + result = f"{{\n 'choices': [\n{choices_str}\n ],\n 'created': {self.created},\n 'model': '{self.model}',\n 'usage': {self.usage}\n}}" return result - def print_verbose(print_statement): if litellm.set_verbose: print(f"LiteLLM: {print_statement}")