mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge cc3ffbe89c into e6c4580a31
This commit is contained in:
commit
07e10534fa
8 changed files with 233 additions and 39 deletions
|
|
@ -0,0 +1,35 @@
|
|||
DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key";
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key";
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key";
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
|
||||
|
||||
UPDATE "LiteLLM_DailyAgentSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
UPDATE "LiteLLM_DailyEndUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
UPDATE "LiteLLM_DailyOrganizationSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
UPDATE "LiteLLM_DailyTagSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
UPDATE "LiteLLM_DailyTeamSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
UPDATE "LiteLLM_DailyUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL;
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
|
@ -762,7 +762,7 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -797,7 +797,7 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([organization_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -831,7 +831,7 @@ model LiteLLM_DailyEndUserSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([end_user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -865,7 +865,7 @@ model LiteLLM_DailyAgentSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([agent_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -901,7 +901,7 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([team_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -937,7 +937,7 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([tag, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -46,7 +46,15 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP
|
|||
# The unique constraint's columns after the entity id, in constraint order. A NULL can
|
||||
# never match itself in a unique index, so every one of these is normalized to '': the
|
||||
# conflict target has to be NULL-free or the row is re-inserted on every single flush.
|
||||
_KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
|
||||
_KEY_COLUMNS: Final = (
|
||||
"date",
|
||||
"api_key",
|
||||
"model",
|
||||
"model_group",
|
||||
"custom_llm_provider",
|
||||
"mcp_namespaced_tool_name",
|
||||
"endpoint",
|
||||
)
|
||||
|
||||
_COUNTER_COLUMNS: Final = (
|
||||
"prompt_tokens",
|
||||
|
|
@ -129,7 +137,6 @@ def _row_params(
|
|||
return (
|
||||
str(uuid.uuid4()),
|
||||
*key,
|
||||
None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),
|
||||
*(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS),
|
||||
*(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS),
|
||||
*((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()),
|
||||
|
|
@ -141,7 +148,6 @@ def _insert_columns(table: DailySpendTable) -> tuple[str, ...]:
|
|||
"id",
|
||||
table.entity_id_column,
|
||||
*_KEY_COLUMNS,
|
||||
"model_group",
|
||||
*_COUNTER_COLUMNS,
|
||||
*_SPEND_COLUMNS,
|
||||
*(("request_id",) if table.carries_request_id else ()),
|
||||
|
|
|
|||
|
|
@ -1495,6 +1495,10 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _daily_transaction_key(*parts: object | None) -> str:
|
||||
return json.dumps(tuple(str(part or "") for part in parts), separators=(",", ":"))
|
||||
|
||||
# fmt: off
|
||||
|
||||
@overload
|
||||
|
|
@ -1940,7 +1944,16 @@ class DBSpendUpdateWriter:
|
|||
return
|
||||
|
||||
endpoint_str: Final = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{payload['user']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key: Final = self._daily_transaction_key(
|
||||
payload["user"],
|
||||
base_daily_transaction["date"],
|
||||
payload["api_key"],
|
||||
payload["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction: Final = DailyUserSpendTransaction(user_id=payload["user"], **base_daily_transaction)
|
||||
await self.daily_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
||||
|
|
@ -1963,7 +1976,16 @@ class DBSpendUpdateWriter:
|
|||
return
|
||||
|
||||
endpoint_str: Final = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key: Final = self._daily_transaction_key(
|
||||
payload["team_id"],
|
||||
base_daily_transaction["date"],
|
||||
payload["api_key"],
|
||||
payload["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction: Final = DailyTeamSpendTransaction(team_id=payload["team_id"], **base_daily_transaction)
|
||||
await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
||||
|
|
@ -1996,7 +2018,16 @@ class DBSpendUpdateWriter:
|
|||
return
|
||||
|
||||
endpoint_str: Final = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key: Final = self._daily_transaction_key(
|
||||
org_id,
|
||||
base_daily_transaction["date"],
|
||||
payload_with_org["api_key"],
|
||||
payload_with_org["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload_with_org["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction: Final = DailyOrganizationSpendTransaction(organization_id=org_id, **base_daily_transaction)
|
||||
await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
||||
|
|
@ -2029,7 +2060,16 @@ class DBSpendUpdateWriter:
|
|||
return
|
||||
|
||||
endpoint_str: Final = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key: Final = self._daily_transaction_key(
|
||||
end_user_id,
|
||||
base_daily_transaction["date"],
|
||||
payload_with_end_user_id["api_key"],
|
||||
payload_with_end_user_id["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload_with_end_user_id["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction: Final = DailyEndUserSpendTransaction(end_user_id=end_user_id, **base_daily_transaction)
|
||||
await self.daily_end_user_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
||||
|
|
@ -2056,7 +2096,16 @@ class DBSpendUpdateWriter:
|
|||
if base_daily_transaction is None:
|
||||
return
|
||||
endpoint_str: Final = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key: Final = self._daily_transaction_key(
|
||||
payload["agent_id"],
|
||||
base_daily_transaction["date"],
|
||||
payload_with_agent_id["api_key"],
|
||||
payload_with_agent_id["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload_with_agent_id["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction: Final = DailyAgentSpendTransaction(agent_id=payload["agent_id"], **base_daily_transaction)
|
||||
await self.daily_agent_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
||||
|
|
@ -2087,7 +2136,16 @@ class DBSpendUpdateWriter:
|
|||
raise ValueError(f"Invalid request_tags: {payload['request_tags']}")
|
||||
for tag in request_tags:
|
||||
endpoint_str = base_daily_transaction.get("endpoint") or ""
|
||||
daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
|
||||
daily_transaction_key = self._daily_transaction_key(
|
||||
tag,
|
||||
base_daily_transaction["date"],
|
||||
payload["api_key"],
|
||||
payload["model"],
|
||||
base_daily_transaction.get("model_group"),
|
||||
payload["custom_llm_provider"],
|
||||
base_daily_transaction.get("mcp_namespaced_tool_name"),
|
||||
endpoint_str,
|
||||
)
|
||||
daily_transaction = DailyTagSpendTransaction(
|
||||
tag=tag, **base_daily_transaction, request_id=payload["request_id"]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -762,7 +762,7 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -797,7 +797,7 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([organization_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -831,7 +831,7 @@ model LiteLLM_DailyEndUserSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([end_user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -865,7 +865,7 @@ model LiteLLM_DailyAgentSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([agent_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -901,7 +901,7 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([team_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -937,7 +937,7 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([tag, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -762,7 +762,7 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -797,7 +797,7 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([organization_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -831,7 +831,7 @@ model LiteLLM_DailyEndUserSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([end_user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -865,7 +865,7 @@ model LiteLLM_DailyAgentSpend {
|
|||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([agent_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -901,7 +901,7 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([team_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -937,7 +937,7 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@unique([tag, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ USER_TABLE = DAILY_SPEND_TABLES["user"]
|
|||
|
||||
# Every nullable member of the unique constraint, so a test that only varied the provider
|
||||
# cannot pass while a sibling column still leaks a NULL into the conflict target.
|
||||
NULLABLE_KEY_COLUMNS = ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
|
||||
NULLABLE_KEY_COLUMNS = ("model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
|
||||
|
||||
|
||||
def tag_txn(**overrides):
|
||||
|
|
@ -78,6 +78,15 @@ def test_distinct_keys_are_not_merged_and_are_ordered_deterministically():
|
|||
assert merged == merge_by_conflict_key(TAG_TABLE, tuple(reversed(unordered)))
|
||||
|
||||
|
||||
def test_distinct_model_groups_are_not_merged():
|
||||
merged = merge_by_conflict_key(
|
||||
TAG_TABLE,
|
||||
(tag_txn(model_group="public-a"), tag_txn(model_group="public-b")),
|
||||
)
|
||||
|
||||
assert len(merged) == 2
|
||||
|
||||
|
||||
def test_one_statement_carries_every_row_in_the_batch():
|
||||
batch = merge_by_conflict_key(TAG_TABLE, tuple(tag_txn(tag=f"team-{i}") for i in range(100)))
|
||||
|
||||
|
|
@ -98,7 +107,8 @@ def test_conflict_target_is_the_full_unique_constraint():
|
|||
conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql)
|
||||
assert conflict_target is not None
|
||||
assert conflict_target.group(1) == (
|
||||
'"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"'
|
||||
'"tag", "date", "api_key", "model", "model_group", "custom_llm_provider", '
|
||||
'"mcp_namespaced_tool_name", "endpoint"'
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import re
|
|||
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from redis.exceptions import DataError
|
||||
|
|
@ -303,6 +303,7 @@ async def test_update_daily_spend_with_null_entity_id():
|
|||
assert _row_values(statement, "date") == ["2024-01-01"]
|
||||
assert _row_values(statement, "api_key") == ["test-api-key"]
|
||||
assert _row_values(statement, "model") == ["gpt-4"]
|
||||
assert _row_values(statement, "model_group") == [""]
|
||||
assert _row_values(statement, "custom_llm_provider") == ["openai"]
|
||||
assert _row_values(statement, "mcp_namespaced_tool_name") == [""]
|
||||
assert _row_values(statement, "endpoint") == [""]
|
||||
|
|
@ -330,6 +331,92 @@ def _daily_txn(user_id: str = "user1") -> dict:
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_user_transaction_key_includes_model_group():
|
||||
writer = DBSpendUpdateWriter()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_request_status = MagicMock(return_value="success")
|
||||
writer.daily_spend_update_queue.add_update = AsyncMock()
|
||||
base_payload = {
|
||||
"request_id": "req-model-group-key",
|
||||
"user": "test-user",
|
||||
"startTime": "2024-01-01T12:00:00",
|
||||
"api_key": "test-key",
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"custom_llm_provider": "openai",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"spend": 0.2,
|
||||
"metadata": '{"usage_object": {}}',
|
||||
}
|
||||
|
||||
await writer.add_spend_log_transaction_to_daily_user_transaction(
|
||||
payload={**base_payload, "model_group": "a"},
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
await writer.add_spend_log_transaction_to_daily_user_transaction(
|
||||
payload={**base_payload, "model_group": "b"},
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
|
||||
queued_keys = [
|
||||
next(iter(call_kwargs.kwargs["update"].keys()))
|
||||
for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list
|
||||
]
|
||||
|
||||
assert queued_keys == [
|
||||
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","",""]',
|
||||
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","b","openai","",""]',
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_user_transaction_key_includes_mcp_tool_name():
|
||||
writer = DBSpendUpdateWriter()
|
||||
writer.daily_spend_update_queue = MagicMock()
|
||||
writer.daily_spend_update_queue.add_update = AsyncMock()
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
base_payload = {
|
||||
"request_id": "req-mcp-tool-key",
|
||||
"user": "test-user",
|
||||
"startTime": "2024-01-01T12:00:00",
|
||||
"api_key": "test-key",
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"model_group": "a",
|
||||
"custom_llm_provider": "openai",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"spend": 0.2,
|
||||
"metadata": '{"usage_object": {}}',
|
||||
}
|
||||
|
||||
await writer.add_spend_log_transaction_to_daily_user_transaction(
|
||||
payload={**base_payload, "mcp_namespaced_tool_name": "server_a.tool"},
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
await writer.add_spend_log_transaction_to_daily_user_transaction(
|
||||
payload={**base_payload, "mcp_namespaced_tool_name": "server_b.tool"},
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
|
||||
queued_keys = [
|
||||
next(iter(call_kwargs.kwargs["update"].keys()))
|
||||
for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list
|
||||
]
|
||||
|
||||
assert queued_keys == [
|
||||
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_a.tool",""]',
|
||||
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_b.tool",""]',
|
||||
]
|
||||
|
||||
|
||||
def test_daily_transaction_key_is_not_delimiter_ambiguous():
|
||||
assert DBSpendUpdateWriter._daily_transaction_key(
|
||||
"user", "model_a", "group"
|
||||
) != DBSpendUpdateWriter._daily_transaction_key("user", "model", "a_group")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors():
|
||||
# Regression for the double-apply hazard: a ReadTimeout means the batch was
|
||||
|
|
@ -626,7 +713,7 @@ async def test_update_tag_db_with_valid_tags():
|
|||
"""
|
||||
Test that _update_tag_db correctly processes valid tags and adds them to the spend update queue.
|
||||
"""
|
||||
from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
|
||||
writer = DBSpendUpdateWriter()
|
||||
mock_prisma = MagicMock()
|
||||
|
|
@ -968,8 +1055,6 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
"metadata": '{"usage_object": {}}',
|
||||
}
|
||||
|
||||
# Mock the add_update method to capture what's being added
|
||||
original_add_update = writer.daily_tag_spend_update_queue.add_update
|
||||
writer.daily_tag_spend_update_queue.add_update = AsyncMock()
|
||||
|
||||
await writer.add_spend_log_transaction_to_daily_tag_transaction(
|
||||
|
|
@ -981,8 +1066,8 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
assert writer.daily_tag_spend_update_queue.add_update.call_count == 2
|
||||
|
||||
# Check that request_id is included in both transactions
|
||||
for call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
|
||||
transaction_dict = call[1]["update"]
|
||||
for update_call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
|
||||
transaction_dict = update_call[1]["update"]
|
||||
# Each transaction should have one key with the format tag_date_api_key_model_provider
|
||||
for key, transaction in transaction_dict.items():
|
||||
assert (
|
||||
|
|
@ -1030,7 +1115,7 @@ async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id
|
|||
update_dict = call_args["update"]
|
||||
assert len(update_dict) == 1
|
||||
for key, transaction in update_dict.items():
|
||||
assert key == f"{org_id}_2024-01-01_test-key_gpt-4_openai_"
|
||||
assert key == f'["{org_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]'
|
||||
assert transaction["organization_id"] == org_id
|
||||
assert transaction["date"] == "2024-01-01"
|
||||
assert transaction["api_key"] == "test-key"
|
||||
|
|
@ -1107,7 +1192,7 @@ async def test_add_spend_log_transaction_to_daily_end_user_transaction_injects_e
|
|||
update_dict = call_args["update"]
|
||||
assert len(update_dict) == 1
|
||||
for key, transaction in update_dict.items():
|
||||
assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai_"
|
||||
assert key == f'["{end_user_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]'
|
||||
assert transaction["end_user_id"] == end_user_id
|
||||
assert transaction["date"] == "2024-01-01"
|
||||
assert transaction["api_key"] == "test-key"
|
||||
|
|
@ -1183,7 +1268,7 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agen
|
|||
update_dict = call_args["update"]
|
||||
assert len(update_dict) == 1
|
||||
for key, transaction in update_dict.items():
|
||||
assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai_"
|
||||
assert key == f'["{agent_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]'
|
||||
assert transaction["agent_id"] == agent_id
|
||||
assert transaction["date"] == "2024-01-01"
|
||||
assert transaction["api_key"] == "test-key"
|
||||
|
|
@ -1304,7 +1389,7 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type():
|
|||
|
||||
for key, transaction in update_dict.items():
|
||||
# Verify endpoint is included in the key
|
||||
assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions"
|
||||
assert key == '["test-user","2024-01-01","test-key","gpt-4","gpt-4-group","openai","","/chat/completions"]'
|
||||
|
||||
# Verify endpoint is set in the transaction
|
||||
assert transaction["endpoint"] == "/chat/completions"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue