test(test_completion.py): cleanup tests

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():

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()

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):
+class EmbeddingResponse(OpenAIObject):
-        d = super().to_dict_recursive()
+    def __init__(self, id=None, choices=None, created=None, model=None, usage=None, stream=False, response_ms=None):
-        return d
+        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):