diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py index 91da1565f7..7e11779c17 100644 --- a/litellm/tests/test_caching.py +++ b/litellm/tests/test_caching.py @@ -127,6 +127,7 @@ def test_embedding_caching(): start_time = time.time() embedding2 = embedding(model="text-embedding-ada-002", input=text_to_embed, caching=True) end_time = time.time() + print(f"embedding2: {embedding2}") print(f"Embedding 2 response time: {end_time - start_time} seconds") litellm.cache = None @@ -136,7 +137,7 @@ def test_embedding_caching(): print(f"embedding2: {embedding2}") pytest.fail("Error occurred: Embedding caching failed") -# test_embedding_caching() +test_embedding_caching() def test_embedding_caching_azure(): diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index e3d36a1a90..acd415bcab 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -51,16 +51,14 @@ def test_completion_claude(): ) # Add any assertions here to check the response print(response) - print(response.response_ms) print(response.usage) print(response.usage.completion_tokens) print(response["usage"]["completion_tokens"]) # print("new cost tracking") - print(response.cost()) except Exception as e: pytest.fail(f"Error occurred: {e}") -# test_completion_claude() +test_completion_claude() # def test_completion_oobabooga(): # try: @@ -1333,7 +1331,6 @@ def test_completion_ai21(): response = completion(model=model_name, messages=messages, max_tokens=100, temperature=0.8) # Add any assertions here to check the response print(response) - print(response.response_ms) except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -1351,7 +1348,6 @@ def test_completion_deep_infra(): ) # Add any assertions here to check the response print(response) - print(response._response_ms) except Exception as e: pytest.fail(f"Error occurred: {e}") test_completion_deep_infra() @@ -1368,7 +1364,6 @@ def test_completion_deep_infra_mistral(): ) # Add any assertions here to check the response print(response) - print(response._response_ms) except Exception as e: pytest.fail(f"Error occurred: {e}") test_completion_deep_infra_mistral() @@ -1381,7 +1376,6 @@ def test_completion_palm(): response = completion(model=model_name, messages=messages, stop=["stop"]) # Add any assertions here to check the response print(response) - print(response.response_ms) except Exception as e: pytest.fail(f"Error occurred: {e}") # test_completion_palm() diff --git a/litellm/utils.py b/litellm/utils.py index 858386f1ff..cc9e4e1c3c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -291,10 +291,6 @@ class ModelResponse(OpenAIObject): created = int(time.time()) else: created = created - if response_ms: - _response_ms = response_ms - else: - _response_ms = None model = model if usage: usage = usage @@ -319,19 +315,33 @@ class ModelResponse(OpenAIObject): def __setitem__(self, key, value): # Allow dictionary-style assignment of attributes setattr(self, key, value) -class EmbeddingResponse(OpenAIObject): - def __init__(self, id=None, choices=None, created=None, model=None, usage=None, stream=False, response_ms=None, **params): - self.object = "list" - if response_ms: - self._response_ms = response_ms - else: - self._response_ms = None - self.data = [] - self.model = model - def to_dict_recursive(self): - d = super().to_dict_recursive() - return d +class EmbeddingResponse(OpenAIObject): + def __init__(self, id=None, choices=None, created=None, model=None, usage=None, stream=False, response_ms=None): + object = "list" + if response_ms: + _response_ms = response_ms + else: + _response_ms = None + data = [] + model = model + super().__init__(id=id, choices=choices, created=created, model=model, object=object, data=data, usage=usage) + + def __contains__(self, key): + # Define custom behavior for the 'in' operator + return hasattr(self, key) + + def get(self, key, default=None): + # Custom .get() method to access attributes with a default value if the attribute doesn't exist + return getattr(self, key, default) + + def __getitem__(self, key): + # Allow dictionary-style access to attributes + return getattr(self, key) + + def __setitem__(self, key, value): + # Allow dictionary-style assignment of attributes + setattr(self, key, value) class TextChoices(OpenAIObject): def __init__(self, finish_reason=None, index=0, text=None, logprobs=None, **params):