forked from phoenix/litellm-mirror
(feat) insert_data to spend table
This commit is contained in:
parent
4821fa9201
commit
4a5f987512
3 changed files with 26 additions and 7 deletions
|
@ -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]
|
||||||
|
|
|
@ -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(
|
||||||
|
|
|
@ -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("{}")
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue