diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 220bf2f91..cf8350022 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -316,7 +316,7 @@ class LiteLLM_UserTable(LiteLLMBase): class LiteLLM_SpendLogs(LiteLLMBase): - id: str + request_id: str call_type: str startTime: Union[str, None] endTime: Union[str, None] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ab1fea463..f24412a84 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1,7 +1,12 @@ from typing import Optional, List, Any, Literal, Union import os, subprocess, hashlib, importlib, asyncio, copy, json, aiohttp, httpx import litellm, backoff -from litellm.proxy._types import UserAPIKeyAuth, DynamoDBArgs, LiteLLM_VerificationToken +from litellm.proxy._types import ( + UserAPIKeyAuth, + DynamoDBArgs, + LiteLLM_VerificationToken, + LiteLLM_SpendLogs, +) from litellm.caching import DualCache from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter @@ -316,7 +321,7 @@ class PrismaClient: self, key: str, value: Any, - table_name: Literal["users", "keys", "config"], + table_name: Literal["users", "keys", "config", "spend"], ): """ Generic implementation of get data @@ -334,6 +339,10 @@ class PrismaClient: response = await self.db.litellm_config.find_first( # type: ignore where={key: value} # type: ignore ) + elif table_name == "spend": + response = await self.db.l.find_first( # type: ignore + where={key: value} # type: ignore + ) return response except Exception as e: asyncio.create_task( @@ -417,7 +426,7 @@ class PrismaClient: on_backoff=on_backoff, # specifying the function to call on backoff ) async def insert_data( - self, data: dict, table_name: Literal["user", "key", "config"] + self, data: dict, table_name: Literal["user", "key", "config", "spend"] ): """ Add a key to the database. If it already exists, do nothing. @@ -473,8 +482,18 @@ class PrismaClient: ) tasks.append(updated_table_row) - await asyncio.gather(*tasks) + elif table_name == "spend": + db_data = self.jsonify_object(data=data) + new_spend_row = await self.db.litellm_spendlogs.upsert( + where={"request_id": data["request_id"]}, + data={ + "create": {**db_data}, # type: ignore + "update": {}, # don't do anything if it already exists + }, + ) + return new_spend_row + except Exception as e: print_verbose(f"LiteLLM Prisma Client Exception: {e}") asyncio.create_task( diff --git a/schema.prisma b/schema.prisma index df2c1d0b4..ed69f67a7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -34,7 +34,7 @@ model LiteLLM_Config { } model LiteLLM_SpendLogs { - id String @unique + request_id String @unique call_type String startTime DateTime // Assuming start_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field @@ -42,7 +42,7 @@ model LiteLLM_SpendLogs { user String @default("") modelParameters Json @default("{}")// Assuming optional_params is a JSON field messages Json @default("[]") - call_cost Float @default(0.0) + spend Float @default(0.0) response Json @default("{}") usage Json @default("{}") metadata Json @default("{}")