fix(proxy_server.py): enable aggregate queries via /spend/keys

litellm/main.py

@@ -31,6 +31,7 @@
     get_llm_provider,
     get_api_key,
     mock_completion_streaming_obj,
+    async_mock_completion_streaming_obj,
     convert_to_model_response_object,
     token_counter,
     Usage,
@@ -307,6 +308,7 @@ def mock_completion(
     messages: List,
     stream: Optional[bool] = False,
     mock_response: str = "This is a mock request",
+    logging=None,
     **kwargs,
 ):
     """
@@ -335,6 +337,15 @@ def mock_completion(
         model_response = ModelResponse(stream=stream)
         if stream is True:
             # don't try to access stream object,
+            if kwargs.get("acompletion", False) == True:
+                return CustomStreamWrapper(
+                    completion_stream=async_mock_completion_streaming_obj(
+                        model_response, mock_response=mock_response, model=model
+                    ),
+                    model=model,
+                    custom_llm_provider="openai",
+                    logging_obj=logging,
+                )
             response = mock_completion_streaming_obj(
                 model_response, mock_response=mock_response, model=model
             )
@@ -717,7 +728,12 @@ def completion(
         )
         if mock_response:
             return mock_completion(
-                model, messages, stream=stream, mock_response=mock_response
+                model,
+                messages,
+                stream=stream,
+                mock_response=mock_response,
+                logging=logging,
+                acompletion=acompletion,
             )
         if custom_llm_provider == "azure":
             # azure configs

litellm/proxy/hooks/parallel_request_limiter.py

@@ -125,7 +125,7 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti
             # ------------
 
             new_val = {
-                "current_requests": current["current_requests"] - 1,
+                "current_requests": max(current["current_requests"] - 1, 0),
                 "current_tpm": current["current_tpm"] + total_tokens,
                 "current_rpm": current["current_rpm"] + 1,
             }
@@ -183,7 +183,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti
                 }
 
                 new_val = {
-                    "current_requests": current["current_requests"] - 1,
+                    "current_requests": max(current["current_requests"] - 1, 0),
                     "current_tpm": current["current_tpm"],
                     "current_rpm": current["current_rpm"],
                 }

litellm/proxy/proxy_server.py

@@ -3015,7 +3015,16 @@ async def info_key_fn(
     tags=["budget & spend Tracking"],
     dependencies=[Depends(user_api_key_auth)],
 )
-async def spend_key_fn():
+async def spend_key_fn(
+    start_date: Optional[str] = fastapi.Query(
+        default=None,
+        description="Time from which to start viewing key spend",
+    ),
+    end_date: Optional[str] = fastapi.Query(
+        default=None,
+        description="Time till which to view key spend",
+    ),
+):
     """
     View all keys created, ordered by spend
 
@@ -3032,9 +3041,41 @@ async def spend_key_fn():
                 f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
             )
 
-        key_info = await prisma_client.get_data(table_name="key", query_type="find_all")
+        if (
+            start_date is not None
+            and isinstance(start_date, str)
+            and end_date is not None
+            and isinstance(end_date, str)
+        ):
+            # Convert the date strings to datetime objects
+            start_date_obj = datetime.strptime(start_date, "%Y-%m-%d")
+            end_date_obj = datetime.strptime(end_date, "%Y-%m-%d")
+
+            # SQL query
+            response = await prisma_client.db.litellm_spendlogs.group_by(
+                by=["api_key", "startTime"],
+                where={
+                    "startTime": {
+                        "gte": start_date_obj,  # Greater than or equal to Start Date
+                        "lte": end_date_obj,  # Less than or equal to End Date
+                    }
+                },
+                sum={
+                    "spend": True,
+                },
+            )
 
-        return key_info
+            # TODO: Execute SQL query and return the results
+
+            return {
+                "message": "This is your SQL query",
+                "response": response,
+            }
+        else:
+            key_info = await prisma_client.get_data(
+                table_name="key", query_type="find_all"
+            )
+            return key_info
 
     except Exception as e:
         raise HTTPException(

litellm/tests/test_parallel_request_limiter.py

@@ -292,6 +292,7 @@ async def test_normal_router_call():
         model="azure-model",
         messages=[{"role": "user", "content": "Hey, how's it going?"}],
         metadata={"user_api_key": _api_key},
+        mock_response="hello",
     )
     await asyncio.sleep(1)  # success is done in a separate thread
     print(f"response: {response}")
@@ -450,6 +451,7 @@ async def test_streaming_router_call():
         messages=[{"role": "user", "content": "Hey, how's it going?"}],
         stream=True,
         metadata={"user_api_key": _api_key},
+        mock_response="hello",
     )
     async for chunk in response:
         continue
@@ -526,6 +528,7 @@ async def test_streaming_router_tpm_limit():
         messages=[{"role": "user", "content": "Write me a paragraph on the moon"}],
         stream=True,
         metadata={"user_api_key": _api_key},
+        mock_response="hello",
     )
     async for chunk in response:
         continue

litellm/utils.py

@@ -1576,7 +1576,7 @@ async def async_success_handler(
                             # only add to cache once we have a complete streaming response
                             litellm.cache.add_cache(result, **kwargs)
                 if isinstance(callback, CustomLogger):  # custom logger class
-                    print_verbose(f"Async success callbacks: CustomLogger")
+                    print_verbose(f"Async success callbacks: {callback}")
                     if self.stream:
                         if "complete_streaming_response" in self.model_call_details:
                             await callback.async_log_success_event(
@@ -8819,6 +8819,14 @@ def mock_completion_streaming_obj(model_response, mock_response, model):
         yield model_response
 
 
+async def async_mock_completion_streaming_obj(model_response, mock_response, model):
+    for i in range(0, len(mock_response), 3):
+        completion_obj = Delta(role="assistant", content=mock_response)
+        model_response.choices[0].delta = completion_obj
+        model_response.choices[0].finish_reason = "stop"
+        yield model_response
+
+
 ########## Reading Config File ############################
 def read_config_args(config_path) -> dict:
     try: