From 11117665e1ee9e9ac9207f68dc17c63dc6d77e4e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 24 Jun 2024 19:32:52 -0700
Subject: [PATCH] test - spend/calculate endpoints

---
 .../spend_management_endpoints.py             |   9 +-
 .../tests/test_spend_calculate_endpoint.py    | 103 ++++++++++++++++++
 2 files changed, 110 insertions(+), 2 deletions(-)
 create mode 100644 litellm/tests/test_spend_calculate_endpoint.py

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 8089c7acbe..abbdc3419e 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -1298,9 +1298,14 @@ async def calculate_spend(request: SpendCalculateRequest):
                 )
             else:
                 _cost = completion_cost(model=request.model, messages=request.messages)
-        else:
-            _completion_response = litellm.ModelResponse(request.completion_response)
+        elif request.completion_response is not None:
+            _completion_response = litellm.ModelResponse(**request.completion_response)
             _cost = completion_cost(completion_response=_completion_response)
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail="Bad Request - Either 'model' or 'completion_response' must be provided",
+            )
         return {"cost": _cost}
     except Exception as e:
         if isinstance(e, HTTPException):
diff --git a/litellm/tests/test_spend_calculate_endpoint.py b/litellm/tests/test_spend_calculate_endpoint.py
new file mode 100644
index 0000000000..f8aff337ec
--- /dev/null
+++ b/litellm/tests/test_spend_calculate_endpoint.py
@@ -0,0 +1,103 @@
+import os
+import sys
+
+import pytest
+from dotenv import load_dotenv
+from fastapi import Request
+from fastapi.routing import APIRoute
+
+import litellm
+from litellm.proxy._types import SpendCalculateRequest
+from litellm.proxy.spend_tracking.spend_management_endpoints import calculate_spend
+from litellm.router import Router
+
+# this file is to test litellm/proxy
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+
+@pytest.mark.asyncio
+async def test_spend_calc_model_messages():
+    cost_obj = await calculate_spend(
+        request=SpendCalculateRequest(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "user", "content": "What is the capital of France?"},
+            ],
+        )
+    )
+
+    print("calculated cost", cost_obj)
+    cost = cost_obj["cost"]
+    assert cost > 0.0
+
+
+@pytest.mark.asyncio
+async def test_spend_calc_model_on_router_messages():
+    from litellm.proxy.proxy_server import llm_router as init_llm_router
+
+    temp_llm_router = Router(
+        model_list=[
+            {
+                "model_name": "special-llama-model",
+                "litellm_params": {
+                    "model": "groq/llama3-8b-8192",
+                },
+            }
+        ]
+    )
+
+    setattr(litellm.proxy.proxy_server, "llm_router", temp_llm_router)
+
+    cost_obj = await calculate_spend(
+        request=SpendCalculateRequest(
+            model="special-llama-model",
+            messages=[
+                {"role": "user", "content": "What is the capital of France?"},
+            ],
+        )
+    )
+
+    print("calculated cost", cost_obj)
+    _cost = cost_obj["cost"]
+
+    assert _cost > 0.0
+
+    # set router to init value
+    setattr(litellm.proxy.proxy_server, "llm_router", init_llm_router)
+
+
+@pytest.mark.asyncio
+async def test_spend_calc_using_response():
+    cost_obj = await calculate_spend(
+        request=SpendCalculateRequest(
+            completion_response={
+                "id": "chatcmpl-3bc7abcd-f70b-48ab-a16c-dfba0b286c86",
+                "choices": [
+                    {
+                        "finish_reason": "stop",
+                        "index": 0,
+                        "message": {
+                            "content": "Yooo! What's good?",
+                            "role": "assistant",
+                        },
+                    }
+                ],
+                "created": "1677652288",
+                "model": "groq/llama3-8b-8192",
+                "object": "chat.completion",
+                "system_fingerprint": "fp_873a560973",
+                "usage": {
+                    "completion_tokens": 8,
+                    "prompt_tokens": 12,
+                    "total_tokens": 20,
+                },
+            }
+        )
+    )
+
+    print("calculated cost", cost_obj)
+    cost = cost_obj["cost"]
+    assert cost > 0.0