diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index c9511669f9a..734fe46bb77 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 0e995cd9c30..5368dafc7ea 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -791,7 +791,8 @@ def test_completion_ai21(): try: response = completion(model=model_name, messages=messages) # Add any assertions here to check the response - print(response) + print(response["response_ms"]) + print(response.response_ms) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index 89b09206025..a974fc8bb3d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -123,7 +123,7 @@ class StreamingChoices(OpenAIObject): self.delta = Delta() class ModelResponse(OpenAIObject): - def __init__(self, id=None, choices=None, created=None, model=None, usage=None, stream=False, **params): + def __init__(self, id=None, choices=None, created=None, model=None, usage=None, stream=False, response_ms=None, **params): if stream: self.object = "chat.completion.chunk" self.choices = [StreamingChoices()] @@ -141,6 +141,10 @@ class ModelResponse(OpenAIObject): self.created = int(time.time()) else: self.created = created + if response_ms: + response_ms = response_ms + else: + response_ms = None self.model = model self.usage = ( usage @@ -624,6 +628,7 @@ def client(original_function): ) # don't interrupt execution of main thread my_thread.start() # RETURN RESULT + result.response_ms = (end_time - start_time).total_seconds() * 1000 # return response latency in ms like openai return result except Exception as e: traceback_exception = traceback.format_exc() diff --git a/pyproject.toml b/pyproject.toml index b8de7bb568e..707d9ac0b6a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.717" +version = "0.1.718" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"