mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge b7193e6acd into ebf829c2fa
This commit is contained in:
commit
73e91735dd
17 changed files with 268 additions and 18 deletions
|
|
@ -0,0 +1,10 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- Records the team that sent each request on its daily tag spend rows, so one tag used by
|
||||
-- several teams lands in one row per team. "" means the request had no team, and every row
|
||||
-- written before this migration keeps "" (attribution is not retroactive).
|
||||
--
|
||||
-- NOT NULL DEFAULT '' rather than nullable: the daily spend upsert normalizes every conflict
|
||||
-- column to '' because a NULL never matches itself in a unique index. On PostgreSQL 11+ adding
|
||||
-- a column with a constant default is a metadata-only change, so this does not rewrite the table.
|
||||
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "team_id" TEXT NOT NULL DEFAULT '';
|
||||
|
|
@ -0,0 +1,14 @@
|
|||
-- CreateIndex (CONCURRENTLY)
|
||||
--
|
||||
-- The new conflict target for the daily tag spend upsert: the old unique key plus team_id.
|
||||
-- Built before the old unique index is dropped (next migrations), so the table is never
|
||||
-- without a unique index the writer can arbitrate on.
|
||||
--
|
||||
-- Disclaimer:
|
||||
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
|
||||
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
|
||||
-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and
|
||||
-- recreated.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key" ON "LiteLLM_DailyTagSpend"("tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
-- CreateIndex (CONCURRENTLY)
|
||||
--
|
||||
-- Serves team-scoped tag reads: WHERE team_id = ANY(...) AND date BETWEEN ... grouped by tag.
|
||||
--
|
||||
-- Disclaimer:
|
||||
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
|
||||
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
|
||||
-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and
|
||||
-- recreated.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_team_id_date_tag_idx" ON "LiteLLM_DailyTagSpend"("team_id", "date", "tag");
|
||||
|
|
@ -0,0 +1,18 @@
|
|||
-- DropIndex (CONCURRENTLY)
|
||||
--
|
||||
-- Drops the unique key without team_id. With it in place, two teams sending the same tag on
|
||||
-- the same day for the same key and model would collide on it, so it has to go once the
|
||||
-- team-aware unique index exists.
|
||||
--
|
||||
-- Upgrade note: an old-version writer still targets this index in its ON CONFLICT clause and
|
||||
-- fails every flush once it is gone. Drain and stop old instances before running migrations.
|
||||
-- An old instance booting against the migrated database also logs "prisma schema out of sync"
|
||||
-- with SQL that drops team_id and both new indexes. Do not run it: it deletes team attribution
|
||||
-- and its CREATE UNIQUE INDEX fails once two teams share a key.
|
||||
--
|
||||
-- Disclaimer:
|
||||
-- - DROP INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
|
||||
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
|
||||
|
|
@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend {
|
|||
id String @id @default(uuid())
|
||||
request_id String?
|
||||
tag String?
|
||||
team_id String @default("") // "" = the request had no team
|
||||
date String
|
||||
api_key String
|
||||
model String?
|
||||
|
|
@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([team_id, date, tag])
|
||||
@@index([api_key])
|
||||
@@index([model])
|
||||
@@index([mcp_namespaced_tool_name])
|
||||
|
|
|
|||
|
|
@ -5223,6 +5223,7 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction):
|
|||
class DailyTagSpendTransaction(BaseDailySpendTransaction):
|
||||
request_id: str | None
|
||||
tag: str
|
||||
team_id: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class DailyAgentSpendTransaction(BaseDailySpendTransaction):
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ class DailyBaselineAttribution(LiteLLMBaseModel):
|
|||
model_group: str | None = None
|
||||
endpoint: str | None = None
|
||||
mcp_namespaced_tool_name: str | None = None
|
||||
team_id: str = ""
|
||||
targets: tuple[DailyBaselineTarget, ...] = ()
|
||||
|
||||
def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow:
|
||||
|
|
@ -71,6 +72,7 @@ class DailyBaselineAttribution(LiteLLMBaseModel):
|
|||
"model_group": self.model_group,
|
||||
"endpoint": self.endpoint,
|
||||
"mcp_namespaced_tool_name": self.mcp_namespaced_tool_name,
|
||||
**({"team_id": self.team_id} if table.team_scoped else {}),
|
||||
table.entity_id_column: target.entity_id,
|
||||
"request_id": request_id,
|
||||
"autorouter_savings_spend": savings_delta,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ class DailySpendTable:
|
|||
name: str
|
||||
entity_id_column: str
|
||||
carries_request_id: bool = False
|
||||
team_scoped: bool = False
|
||||
|
||||
|
||||
DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType(
|
||||
|
|
@ -41,7 +42,9 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP
|
|||
"org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"),
|
||||
"end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"),
|
||||
"agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"),
|
||||
"tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True),
|
||||
"tag": DailySpendTable(
|
||||
name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, team_scoped=True
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -124,9 +127,14 @@ def _as_float(value: object) -> float:
|
|||
return float(value) if isinstance(value, (int, float)) else 0.0
|
||||
|
||||
|
||||
def _conflict_columns(table: DailySpendTable) -> tuple[str, ...]:
|
||||
"""The table's unique constraint, in constraint order."""
|
||||
return (table.entity_id_column, *(("team_id",) if table.team_scoped else ()), *_KEY_COLUMNS)
|
||||
|
||||
|
||||
def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]:
|
||||
"""The tuple the database arbitrates the upsert on, normalized free of NULLs."""
|
||||
return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS))
|
||||
return tuple(_as_text(transaction.get(column)) for column in _conflict_columns(table))
|
||||
|
||||
|
||||
def _merge(group: Sequence[SpendRow]) -> SpendRow:
|
||||
|
|
@ -174,8 +182,7 @@ def _row_params(
|
|||
def _insert_columns(table: DailySpendTable) -> tuple[str, ...]:
|
||||
return (
|
||||
"id",
|
||||
table.entity_id_column,
|
||||
*_KEY_COLUMNS,
|
||||
*_conflict_columns(table),
|
||||
"model_group",
|
||||
*_COUNTER_COLUMNS,
|
||||
*_SPEND_COLUMNS,
|
||||
|
|
@ -213,7 +220,7 @@ def build_bulk_upsert(
|
|||
sql: Final = (
|
||||
f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n'
|
||||
f"VALUES {rows}\n"
|
||||
f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n"
|
||||
f"ON CONFLICT ({_quoted(_conflict_columns(table))}) DO UPDATE SET\n"
|
||||
f" {increments}{request_id_update},\n"
|
||||
f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1019,6 +1019,7 @@ class DBSpendUpdateWriter:
|
|||
model_group=base.get("model_group"),
|
||||
endpoint=base.get("endpoint"),
|
||||
mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"),
|
||||
team_id=payload.get("team_id") or "", # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults
|
||||
targets=tuple(
|
||||
DailyBaselineTarget(entity=entity, entity_id=identity)
|
||||
for entity, values in bases
|
||||
|
|
@ -3075,13 +3076,25 @@ class DBSpendUpdateWriter:
|
|||
return
|
||||
|
||||
request_tags: Final = daily_spend_entity_ids(payload, "tag")
|
||||
# .get, not indexing: the caller swallows exceptions, so a KeyError would silently drop tag spend
|
||||
team_id: Final = payload.get("team_id") or "" # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults
|
||||
for tag in request_tags:
|
||||
if tag is None:
|
||||
continue
|
||||
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 = json.dumps(
|
||||
[
|
||||
tag,
|
||||
team_id,
|
||||
base_daily_transaction["date"],
|
||||
payload["api_key"],
|
||||
payload["model"],
|
||||
payload["custom_llm_provider"],
|
||||
endpoint_str,
|
||||
]
|
||||
)
|
||||
daily_transaction = DailyTagSpendTransaction(
|
||||
tag=tag, **base_daily_transaction, request_id=payload["request_id"]
|
||||
tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"]
|
||||
)
|
||||
|
||||
await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
|
||||
|
|
|
|||
|
|
@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend {
|
|||
id String @id @default(uuid())
|
||||
request_id String?
|
||||
tag String?
|
||||
team_id String @default("") // "" = the request had no team
|
||||
date String
|
||||
api_key String
|
||||
model String?
|
||||
|
|
@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([team_id, date, tag])
|
||||
@@index([api_key])
|
||||
@@index([model])
|
||||
@@index([mcp_namespaced_tool_name])
|
||||
|
|
|
|||
|
|
@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend {
|
|||
id String @id @default(uuid())
|
||||
request_id String?
|
||||
tag String?
|
||||
team_id String @default("") // "" = the request had no team
|
||||
date String
|
||||
api_key String
|
||||
model String?
|
||||
|
|
@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([team_id, date, tag])
|
||||
@@index([api_key])
|
||||
@@index([model])
|
||||
@@index([mcp_namespaced_tool_name])
|
||||
|
|
|
|||
|
|
@ -15,10 +15,14 @@ Schema = Mapping[str, Mapping[str, str]]
|
|||
MODELS: Final = re.compile(r"^[ \t]*model\s+(\w+)\s*\{(.*?)^[ \t]*\}", re.MULTILINE | re.DOTALL)
|
||||
|
||||
|
||||
def _is_field(part: list[str]) -> bool:
|
||||
return len(part) >= 2 and not part[0].startswith(("@@", "//"))
|
||||
|
||||
|
||||
def parse_model_fields(schema: str) -> Schema:
|
||||
def fields(body: str) -> Mapping[str, str]:
|
||||
parts: Final = tuple(line.split() for line in body.splitlines())
|
||||
return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if len(part) >= 2})
|
||||
return MappingProxyType({part[0]: " ".join(part[1:]) for part in parts if _is_field(part)})
|
||||
|
||||
return MappingProxyType({match[1]: fields(match[2]) for match in MODELS.finditer(schema)})
|
||||
|
||||
|
|
|
|||
45
tests/integration/spend/test_team_tag_spend_attribution.py
Normal file
45
tests/integration/spend/test_team_tag_spend_attribution.py
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
import uuid
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
|
||||
_SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002
|
||||
|
||||
|
||||
def test_one_key_moved_between_teams_records_its_tag_once_per_team(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
tag: Final = f"cost-center:{uuid.uuid4().hex}"
|
||||
teams: Final = tuple(scenario.team(models=[model]) for _ in range(2))
|
||||
calls_per_team: Final = {teams[0]: 1, teams[1]: 2}
|
||||
key: Final = scenario.key(team_id=teams[0], models=[model])
|
||||
for team, calls in calls_per_team.items():
|
||||
gateway.post("/key/update", {"key": key, "team_id": team})
|
||||
for _ in range(calls):
|
||||
reply = gateway.post(
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tagged {uuid.uuid4().hex}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert reply["usage"]["total_tokens"] == 40, reply
|
||||
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT team_id, spend, api_requests FROM "LiteLLM_DailyTagSpend" WHERE tag = %s ORDER BY team_id',
|
||||
(tag,),
|
||||
),
|
||||
lambda values: sum(int(row["api_requests"]) for row in values) == 3,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
assert sorted((row["team_id"], int(row["api_requests"])) for row in rows) == sorted(calls_per_team.items()), (
|
||||
rows
|
||||
)
|
||||
for row in rows:
|
||||
assert float(row["spend"]) == pytest.approx(calls_per_team[row["team_id"]] * _SPEND_PER_CALL), rows
|
||||
|
|
@ -435,7 +435,7 @@ class TestMigrationSQLIdempotency:
|
|||
for migration_name, sql in _guarded_migrations(all_migrations):
|
||||
for line_num, line in _statements(sql):
|
||||
if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search(
|
||||
r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE
|
||||
r"DROP\s+INDEX\s+(?:CONCURRENTLY\s+)?IF\s+EXISTS", line, re.IGNORECASE
|
||||
):
|
||||
violations.append(f" {migration_name}:{line_num}: {line.strip()}")
|
||||
assert (
|
||||
|
|
|
|||
|
|
@ -3,11 +3,13 @@
|
|||
import re
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.daily_spend_bulk_upsert import (
|
||||
DAILY_SPEND_TABLES,
|
||||
DailySpendEntity,
|
||||
build_bulk_upsert,
|
||||
conflict_key,
|
||||
merge_by_conflict_key,
|
||||
|
|
@ -88,10 +90,8 @@ def test_one_statement_carries_every_row_in_the_batch():
|
|||
|
||||
assert sql.count("INSERT INTO") == 1
|
||||
assert len(re.findall(r"ON CONFLICT", sql)) == 1
|
||||
# 25 bound columns per row plus the inlined updated_at, so the row count is what
|
||||
# separates one multi-row statement from a hundred single-row ones.
|
||||
assert len(params) == 100 * 25
|
||||
assert "$2500::text" in sql
|
||||
assert len(params) == 100 * 26
|
||||
assert "$2600::text" in sql
|
||||
assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1
|
||||
|
||||
|
||||
|
|
@ -101,7 +101,7 @@ 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", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"'
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -144,6 +144,41 @@ def test_non_tag_tables_carry_no_request_id_column():
|
|||
assert '"user_id"' in sql
|
||||
|
||||
|
||||
def test_one_tag_from_two_teams_stays_two_rows() -> None:
|
||||
merged: Final = merge_by_conflict_key(
|
||||
TAG_TABLE, (tag_txn(team_id="team-b", spend=2.0), tag_txn(team_id="team-a", spend=1.0))
|
||||
)
|
||||
|
||||
assert [(txn["team_id"], txn["spend"]) for _, txn in merged] == [("team-a", 1.0), ("team-b", 2.0)], merged
|
||||
|
||||
|
||||
def test_an_entry_without_team_id_folds_into_the_unattributed_row() -> None:
|
||||
legacy: Final = tag_txn(spend=0.25)
|
||||
unattributed: Final = tag_txn(team_id="", spend=0.75)
|
||||
|
||||
merged: Final = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed))
|
||||
|
||||
assert len(merged) == 1, merged
|
||||
key, folded = merged[0]
|
||||
assert key[:2] == ("team-a", ""), key
|
||||
assert folded["spend"] == pytest.approx(1.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"])
|
||||
def test_only_the_tag_table_is_keyed_by_team(entity: DailySpendEntity) -> None:
|
||||
table: Final = DAILY_SPEND_TABLES[entity]
|
||||
row: Final = {**tag_txn(), table.entity_id_column: "id-1"}
|
||||
|
||||
sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,)))
|
||||
|
||||
conflict_target: Final = re.search(r"ON CONFLICT \(([^)]*)\)", sql)
|
||||
assert conflict_target is not None
|
||||
assert conflict_target.group(1) == (
|
||||
f'"{table.entity_id_column}", "date", "api_key", "model", "custom_llm_provider", '
|
||||
'"mcp_namespaced_tool_name", "endpoint"'
|
||||
)
|
||||
|
||||
|
||||
class _RecordingDb:
|
||||
def __init__(self) -> None:
|
||||
self.statements: list[tuple[str, tuple[object, ...]]] = []
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceTypes
|
||||
from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS
|
||||
from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem
|
||||
from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendLogsPayload, SpendUpdateQueueItem
|
||||
from litellm.proxy.db.db_spend_update_writer import (
|
||||
_TEAM_ADVISORY_LOCK_SQL,
|
||||
_TEAM_MEMBER_SPEND_SQL,
|
||||
|
|
@ -1493,6 +1493,67 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
)
|
||||
|
||||
|
||||
def _tagged_spend_log(tag: str, team_id: str | None, spend: float) -> SpendLogsPayload:
|
||||
return cast(
|
||||
SpendLogsPayload,
|
||||
{
|
||||
"request_id": f"req-{tag}-{team_id}",
|
||||
"request_tags": json.dumps([tag]),
|
||||
"team_id": team_id,
|
||||
"user": "user-1",
|
||||
"startTime": "2026-10-17T00:00:00",
|
||||
"api_key": "key-1",
|
||||
"model": "gpt-4o",
|
||||
"custom_llm_provider": "openai",
|
||||
"model_group": "gpt-4o",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"spend": spend,
|
||||
"metadata": '{"usage_object": {}}',
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("requests", "expected"),
|
||||
[
|
||||
(
|
||||
(
|
||||
("cost-center:cc-42", "team-a", 1.0),
|
||||
("cost-center:cc-42", "team-b", 2.0),
|
||||
("cost-center:cc-42", None, 0.5),
|
||||
),
|
||||
[
|
||||
("cost-center:cc-42", "", 0.5),
|
||||
("cost-center:cc-42", "team-a", 1.0),
|
||||
("cost-center:cc-42", "team-b", 2.0),
|
||||
],
|
||||
),
|
||||
(
|
||||
(("cost_center", "west", 1.0), ("cost", "center_west", 2.0)),
|
||||
[("cost", "center_west", 2.0), ("cost_center", "west", 1.0)],
|
||||
),
|
||||
],
|
||||
ids=["one-tag-two-teams", "underscore-joined-names-collide"],
|
||||
)
|
||||
async def test_each_tag_and_team_pair_is_queued_as_its_own_entry(
|
||||
requests: tuple[tuple[str, str | None, float], ...],
|
||||
expected: list[tuple[str, str, float]],
|
||||
) -> None:
|
||||
writer: Final = DBSpendUpdateWriter()
|
||||
prisma: Final = MagicMock()
|
||||
prisma.get_request_status = MagicMock(return_value="success")
|
||||
|
||||
for tag, team_id, spend in requests:
|
||||
await writer.add_spend_log_transaction_to_daily_tag_transaction(
|
||||
payload=_tagged_spend_log(tag, team_id, spend), prisma_client=prisma
|
||||
)
|
||||
|
||||
queued: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
|
||||
assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == expected, queued
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id_and_queues_update():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -16,6 +16,28 @@ def test_parse_model_fields_retains_types_and_attributes() -> None:
|
|||
assert parse_model_fields(ORIGINAL) == {"Item": {"id": "String @id", "label": "String?"}}
|
||||
|
||||
|
||||
def test_block_attributes_and_comments_are_not_fields() -> None:
|
||||
schema: Final = (
|
||||
"model Item {\n id String @id\n label String?\n // free text\n"
|
||||
' @@unique([label, id], map: "item_label_id_key")\n @@index([label, id])\n @@index([id])\n}\n'
|
||||
)
|
||||
assert parse_model_fields(schema) == {"Item": {"id": "String @id", "label": "String?"}}
|
||||
|
||||
|
||||
def test_changing_a_composite_unique_key_is_not_a_field_change() -> None:
|
||||
before: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id])\n}\n"
|
||||
after: Final = "model Item {\n id String @id\n label String?\n @@unique([label, id, tenant])\n}\n"
|
||||
assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == ()
|
||||
|
||||
|
||||
def test_a_narrowed_column_is_still_caught_next_to_block_attributes() -> None:
|
||||
before: Final = "model Item {\n id String @id\n label String?\n @@index([label, id])\n}\n"
|
||||
after: Final = "model Item {\n id String @id\n label String\n @@index([label, id])\n}\n"
|
||||
assert check_breaking_changes(parse_model_fields(before), parse_model_fields(after)) == (
|
||||
"Breaking: Field Item.label changed from optional to required",
|
||||
)
|
||||
|
||||
|
||||
def test_parse_model_fields_keeps_json_defaults_and_empty_models() -> None:
|
||||
schema: Final = 'model Item {\n data Json @default("{}")\n}\nmodel Empty {\n}\n'
|
||||
assert parse_model_fields(schema) == {"Item": {"data": 'Json @default("{}")'}, "Empty": {}}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue