From 11e3ee44113f9665c7cbe7ba52f8155d1941bc4b Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 22 Jan 2024 20:20:17 -0800
Subject: [PATCH] test: fix tests

---
 litellm/proxy/proxy_server.py               |  2 +-
 litellm/tests/test_key_generate_dynamodb.py |  6 ++++-
 litellm/tests/test_key_generate_prisma.py   |  2 ++
 litellm/utils.py                            | 26 +++++++++++++--------
 4 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index dfb4b70f88..e101aac32f 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -589,7 +589,7 @@ async def track_cost_callback(
                 )
         else:
             raise Exception(
-                f"Model={kwargs['model']} not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
+                f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
             )
     except Exception as e:
         verbose_proxy_logger.debug(f"error in tracking cost callback - {str(e)}")
diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 2cfa9c9531..c6936a99e3 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -179,9 +179,11 @@ def test_call_with_key_over_budget(custom_db_client):
     # 5. Make a call with a key over budget, expect to fail
     setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
-    from litellm._logging import verbose_proxy_logger
+    from litellm._logging import verbose_proxy_logger, verbose_logger
     import logging
 
+    litellm.set_verbose = True
+    verbose_logger.setLevel(logging.DEBUG)
     verbose_proxy_logger.setLevel(logging.DEBUG)
     try:
 
@@ -229,6 +231,7 @@ def test_call_with_key_over_budget(custom_db_client):
                             "user_api_key_user_id": user_id,
                         }
                     },
+                    "response_cost": 0.00002,
                 },
                 completion_response=resp,
             )
@@ -301,6 +304,7 @@ def test_call_with_key_over_budget_stream(custom_db_client):
                             "user_api_key_user_id": user_id,
                         }
                     },
+                    "response_cost": 0.00002,
                 },
                 completion_response=ModelResponse(),
             )
diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index 5ecfc89d75..f03356700e 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -256,6 +256,7 @@ def test_call_with_key_over_budget(prisma_client):
                             "user_api_key_user_id": user_id,
                         }
                     },
+                    "response_cost": 0.00002,
                 },
                 completion_response=resp,
                 start_time=datetime.now(),
@@ -331,6 +332,7 @@ def test_call_with_key_over_budget_stream(prisma_client):
                             "user_api_key_user_id": user_id,
                         }
                     },
+                    "response_cost": 0.00002,
                 },
                 completion_response=ModelResponse(),
                 start_time=datetime.now(),
diff --git a/litellm/utils.py b/litellm/utils.py
index c67b6d39f5..21972373a5 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1066,17 +1066,23 @@ class Logging:
             self.model_call_details["cache_hit"] = cache_hit
             ## if model in model cost map - log the response cost
             ## else set cost to None
-            if (
-                result is not None
-                and (
-                    isinstance(result, ModelResponse)
-                    or isinstance(result, EmbeddingResponse)
-                )
-                and result.model in litellm.model_cost
+            verbose_logger.debug(f"Model={self.model}; result={result}")
+            if result is not None and (
+                isinstance(result, ModelResponse)
+                or isinstance(result, EmbeddingResponse)
             ):
-                self.model_call_details["response_cost"] = litellm.completion_cost(
-                    completion_response=result,
-                )
+                try:
+                    self.model_call_details["response_cost"] = litellm.completion_cost(
+                        completion_response=result,
+                    )
+                    verbose_logger.debug(
+                        f"Model={self.model}; cost={self.model_call_details['response_cost']}"
+                    )
+                except litellm.NotFoundError as e:
+                    verbose_logger.debug(
+                        f"Model={self.model} not found in completion cost map."
+                    )
+                    self.model_call_details["response_cost"] = None
             else:  # streaming chunks + image gen.
                 self.model_call_details["response_cost"] = None