(feat) insert_data to spend table

This commit is contained in:
ishaan-jaff 2024-01-18 10:09:02 -08:00
parent 4821fa9201
commit 4a5f987512
3 changed files with 26 additions and 7 deletions

View file

@ -316,7 +316,7 @@ class LiteLLM_UserTable(LiteLLMBase):
class LiteLLM_SpendLogs(LiteLLMBase): class LiteLLM_SpendLogs(LiteLLMBase):
id: str request_id: str
call_type: str call_type: str
startTime: Union[str, None] startTime: Union[str, None]
endTime: Union[str, None] endTime: Union[str, None]

View file

@ -1,7 +1,12 @@
from typing import Optional, List, Any, Literal, Union from typing import Optional, List, Any, Literal, Union
import os, subprocess, hashlib, importlib, asyncio, copy, json, aiohttp, httpx import os, subprocess, hashlib, importlib, asyncio, copy, json, aiohttp, httpx
import litellm, backoff 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.caching import DualCache
from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler
from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter from litellm.proxy.hooks.max_budget_limiter import MaxBudgetLimiter
@ -316,7 +321,7 @@ class PrismaClient:
self, self,
key: str, key: str,
value: Any, value: Any,
table_name: Literal["users", "keys", "config"], table_name: Literal["users", "keys", "config", "spend"],
): ):
""" """
Generic implementation of get data Generic implementation of get data
@ -334,6 +339,10 @@ class PrismaClient:
response = await self.db.litellm_config.find_first( # type: ignore response = await self.db.litellm_config.find_first( # type: ignore
where={key: value} # 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 return response
except Exception as e: except Exception as e:
asyncio.create_task( asyncio.create_task(
@ -417,7 +426,7 @@ class PrismaClient:
on_backoff=on_backoff, # specifying the function to call on backoff on_backoff=on_backoff, # specifying the function to call on backoff
) )
async def insert_data( 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. Add a key to the database. If it already exists, do nothing.
@ -473,8 +482,18 @@ class PrismaClient:
) )
tasks.append(updated_table_row) tasks.append(updated_table_row)
await asyncio.gather(*tasks) 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: except Exception as e:
print_verbose(f"LiteLLM Prisma Client Exception: {e}") print_verbose(f"LiteLLM Prisma Client Exception: {e}")
asyncio.create_task( asyncio.create_task(

View file

@ -34,7 +34,7 @@ model LiteLLM_Config {
} }
model LiteLLM_SpendLogs { model LiteLLM_SpendLogs {
id String @unique request_id String @unique
call_type String call_type String
startTime DateTime // Assuming start_time is a DateTime field startTime DateTime // Assuming start_time is a DateTime field
endTime DateTime // Assuming end_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field
@ -42,7 +42,7 @@ model LiteLLM_SpendLogs {
user String @default("") user String @default("")
modelParameters Json @default("{}")// Assuming optional_params is a JSON field modelParameters Json @default("{}")// Assuming optional_params is a JSON field
messages Json @default("[]") messages Json @default("[]")
call_cost Float @default(0.0) spend Float @default(0.0)
response Json @default("{}") response Json @default("{}")
usage Json @default("{}") usage Json @default("{}")
metadata Json @default("{}") metadata Json @default("{}")