scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Merge pull request #46313 from BerriAI/litellm_backport_44243_44508_stable_1_105_x

Browse source 
fix(proxy): backport #44243 and #44508 to stable/1.105.x for v1.105.1
This commit is contained in:
Kerry Lu 2026-10-13 15:48:16 -07:00 • committed by GitHub
commit d665dc1a63
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 646 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/litellm_core_utils/litellm_logging.py
View file

@ -718,6 +718,7 @@ class Logging(LiteLLMLoggingBaseClass):
# enqueue closure here instead of firing it immediately.
self._defer_async_logging: bool = False
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._async_success_scheduled: bool = False
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
@ -2256,6 +2257,15 @@ class Logging(LiteLLMLoggingBaseClass):
except Exception:
return True
def claim_async_success_log(self) -> bool:
"""One async success log per request. The innermost @client wrapper always exits first,
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
False and schedules nothing."""
if self._async_success_scheduled:
return False
self._async_success_scheduled = True
return True
def has_run_logging(
self,
event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"],

25
litellm/proxy/db/db_spend_update_writer.py
View file

@ -70,6 +70,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendUpdateQueue,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
from litellm.proxy.spend_tracking.compression_savings import (
extract_compression_saved_tokens,
@ -728,6 +729,20 @@ class DBSpendUpdateWriter:
except Exception as e:
verbose_proxy_logger.debug("_enqueue_tool_usage_transaction error (non-blocking): %s", e)
async def _enqueue_model_usage_transaction(
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient,
) -> None:
try:
transaction: Final = build_model_usage_transaction(payload)
if transaction is None:
return
async with prisma_client._model_usage_transactions_lock:
prisma_client.model_usage_transactions.append(transaction)
except Exception as e:
verbose_proxy_logger.debug("_enqueue_model_usage_transaction error (non-blocking): %s", e)
async def _enqueue_autorouter_turn_transaction(
self,
payload: SpendLogsPayload,
@ -1144,15 +1159,7 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
from litellm.proxy.db.model_usage_rollup import increment_daily_model_usage
await increment_daily_model_usage(prisma_client=prisma_client, payload=payload_copy)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: increment_daily_model_usage failed: %s",
traceback.format_exc(),
)
await self._enqueue_model_usage_transaction(payload=payload_copy, prisma_client=prisma_client)
async def _update_key_db(
self,

165
litellm/proxy/db/model_usage_rollup.py
View file

@ -1,5 +1,12 @@
from __future__ import annotations
import asyncio
import random
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import Final
from itertools import groupby
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter, ValidationError
@ -8,15 +15,48 @@ from litellm.constants import (
MODEL_INSIGHTS_DEFAULT_TASK,
MODEL_INSIGHTS_TASK_TAG_PREFIX,
)
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES, SpendLogsPayload
from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import DailyModelUsageRepository
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_METADATA: Final = TypeAdapter(dict[str, object])
_TAGS: Final = TypeAdapter(list[object])
class _UpsertTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
class _ModelUsageBatch(Protocol):
litellm_dailymodelusage: _UpsertTable
class _ModelUsageBatchManager(Protocol):
async def __aenter__(self) -> _ModelUsageBatch: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
@dataclass(frozen=True, slots=True)
class ModelUsageKey:
date: str
model_group: str
model: str
custom_llm_provider: str
task_type: str
@dataclass(frozen=True, slots=True)
class ModelUsageTransaction:
key: ModelUsageKey
spend: float
prompt_tokens: int
completion_tokens: int
successful: bool
def model_usage_task_type(request_tags: str) -> str:
try:
tags: Final = _TAGS.validate_json(request_tags)
@ -48,43 +88,88 @@ def _date_from_start_time(start_time: datetime | str) -> str | None:
return start_time[:10] if len(start_time) >= 10 else None
async def increment_daily_model_usage(prisma_client: PrismaClient, payload: SpendLogsPayload) -> None:
def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None:
date: Final = _date_from_start_time(payload["startTime"])
if date is None or _is_internal_call(payload["metadata"]):
return
return None
model: Final = payload["model"] or "unknown"
model_group: Final = payload["model_group"] or model
provider: Final = payload["custom_llm_provider"] or "unknown"
task_type: Final = model_usage_task_type(payload["request_tags"])
successful: Final = 1 if payload["status"] == "success" else 0
failed: Final = 1 - successful
key: Final = {
"date": date,
"model_group": model_group,
"model": model,
"custom_llm_provider": provider,
"task_type": task_type,
}
await DailyModelUsageRepository(prisma_client).table.upsert(
where={"date_model_group_model_custom_llm_provider_task_type": key},
data={
"create": {
**key,
"spend": payload["spend"],
"prompt_tokens": payload["prompt_tokens"],
"completion_tokens": payload["completion_tokens"],
"request_count": 1,
"successful_requests": successful,
"failed_requests": failed,
},
"update": {
"spend": {"increment": payload["spend"]},
"prompt_tokens": {"increment": payload["prompt_tokens"]},
"completion_tokens": {"increment": payload["completion_tokens"]},
"request_count": {"increment": 1},
"successful_requests": {"increment": successful},
"failed_requests": {"increment": failed},
},
},
return ModelUsageTransaction(
key=ModelUsageKey(
date=date,
model_group=payload["model_group"] or model,
model=model,
custom_llm_provider=payload["custom_llm_provider"] or "unknown",
task_type=model_usage_task_type(payload["request_tags"]),
),
spend=payload["spend"],
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
successful=payload["status"] == "success",
)
def _model_usage_batch(prisma_client: PrismaClient) -> _ModelUsageBatchManager:
batch: Final[_ModelUsageBatchManager] = prisma_client.db.batch_()
return batch
def _sort_key(transaction: ModelUsageTransaction) -> tuple[str, str, str, str, str]:
key: Final = transaction.key
return (key.date, key.model_group, key.model, key.custom_llm_provider, key.task_type)
async def flush_model_usage_transactions(
prisma_client: PrismaClient,
transactions: Sequence[ModelUsageTransaction],
n_retry_times: int = 3,
) -> None:
"""One upsert per rollup row for the whole drained batch, in a single transaction and in key order so
concurrent pods take row locks in the same order. Only ConnectError is retried: it proves nothing reached
the database, while a retry after an ambiguous post-send failure could double-count the increments."""
if not transactions:
return
ordered: Final = sorted(transactions, key=_sort_key)
for attempt in range(n_retry_times + 1):
try:
async with _model_usage_batch(prisma_client) as batcher:
for key, grouped in groupby(ordered, key=lambda transaction: transaction.key):
entries = tuple(grouped)
spend = sum(entry.spend for entry in entries)
prompt_tokens = sum(entry.prompt_tokens for entry in entries)
completion_tokens = sum(entry.completion_tokens for entry in entries)
successful = sum(1 for entry in entries if entry.successful)
failed = len(entries) - successful
key_fields = {
"date": key.date,
"model_group": key.model_group,
"model": key.model,
"custom_llm_provider": key.custom_llm_provider,
"task_type": key.task_type,
}
batcher.litellm_dailymodelusage.upsert(
where={"date_model_group_model_custom_llm_provider_task_type": key_fields},
data={
"create": {
**key_fields,
"spend": spend,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"request_count": len(entries),
"successful_requests": successful,
"failed_requests": failed,
},
"update": {
"spend": {"increment": spend},
"prompt_tokens": {"increment": prompt_tokens},
"completion_tokens": {"increment": completion_tokens},
"request_count": {"increment": len(entries)},
"successful_requests": {"increment": successful},
"failed_requests": {"increment": failed},
},
},
)
return
except DB_RETRY_SAFE_ERROR_TYPES:
if attempt >= n_retry_times:
raise
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))

20
litellm/proxy/utils.py
View file

@ -266,6 +266,7 @@ if TYPE_CHECKING:
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord
from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
@ -4400,6 +4401,8 @@ class PrismaClient:
spend_log_write_lock = asyncio.Lock()
tool_usage_transactions: list["ToolUsageTransaction"] = []
_tool_usage_transactions_lock = asyncio.Lock()
model_usage_transactions: ClassVar[list["ModelUsageTransaction"]] = []
_model_usage_transactions_lock = asyncio.Lock()
autorouter_turn_transactions: ClassVar[list["AutoRouterTurnTransaction"]] = []
_autorouter_turn_transactions_lock = asyncio.Lock()
@ -7508,6 +7511,8 @@ async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
spend_queue_size: Final = len(prisma_client.spend_log_transactions)
async with prisma_client._tool_usage_transactions_lock:
tool_queue_size: Final = len(prisma_client.tool_usage_transactions)
async with prisma_client._model_usage_transactions_lock:
model_usage_queue_size: Final = len(prisma_client.model_usage_transactions)
async with prisma_client._autorouter_turn_transactions_lock:
autorouter_queue_size: Final = len(prisma_client.autorouter_turn_transactions)
from litellm.proxy.db.shadow_eval_funnel import pending_shadow_eval_funnel_events
@ -7517,6 +7522,7 @@ async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
return (
spend_queue_size
+ tool_queue_size
+ model_usage_queue_size
+ autorouter_queue_size
+ baseline_queue_size
+ pending_shadow_eval_funnel_events()
@ -7652,6 +7658,20 @@ async def _run_spend_logs_job(
tool_tracking_err,
)
async with prisma_client._model_usage_transactions_lock:
model_usage_to_process: Final = prisma_client.model_usage_transactions
prisma_client.model_usage_transactions = []
try:
from litellm.proxy.db.model_usage_rollup import flush_model_usage_transactions
await flush_model_usage_transactions(prisma_client=prisma_client, transactions=model_usage_to_process)
except Exception as model_usage_err:
verbose_proxy_logger.error(
"Spend tracking - model usage flush failed; %s model usage transactions dropped: %s",
len(model_usage_to_process),
model_usage_err,
)
await flush_baseline_accounting(prisma_client)
async with prisma_client._autorouter_turn_transactions_lock:

11
litellm/utils.py
View file

@ -1380,11 +1380,14 @@ def _schedule_async_success_logging(
Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses
bridge) each exit through here with the same logging object and their own shape of the same
response. The immediate path already logs one request once, since the first task marks
``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same
first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log
reads usage from, and a later wrapper never swaps in its client-shaped translation.
response. The innermost wrapper always exits first, since the outer one is awaiting it, so it
claims the async success log here, synchronously, and the outer wrapper returns without
enqueuing: one request queues one handler, whatever the handler tasks do later. The deferred
slot keeps the same innermost-wins rule: the provider-shaped result is the one the spend log
reads usage from, and the outer wrapper never swaps in its client-shaped translation.
"""
if not logging_obj.claim_async_success_log():
return
def _enqueue_async_logging() -> None:
asyncio.create_task(

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.105.0"
version = "1.105.1"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true }
members = ["enterprise", "litellm-proxy-extras"]
[tool.commitizen]
version = "1.105.0"
version = "1.105.1"
version_files = [
"pyproject.toml:^version",
]

111
tests/integration/spend/test_responses_bridge_spend_once.py Normal file
View file

@ -0,0 +1,111 @@
import json
from base64 import b64decode
from collections.abc import Callable
from hashlib import sha256
from typing import Final
from uuid import uuid4
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
PROMPT_TOKENS: Final = 30
COMPLETION_TOKENS: Final = 5
INPUT_RATE: Final = 0.001
OUTPUT_RATE: Final = 0.002
EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE
def _responses_reply(prompt: str, response_id: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return Reply(
body=json.dumps({"object": "list", "data": [{"id": "gpt-6.1-sol", "object": "model"}]}).encode()
)
assert request.method == "POST", request.method
assert request.target == "/v1/responses", request.target
body: Final = object_value(json.loads(request.body))
encoded_input: Final = json.dumps(body["input"])
assert prompt in encoded_input, encoded_input
assert body.get("stream") is not True, body
response: Final = {
"id": response_id,
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": body["model"],
"output": [
{
"type": "message",
"id": f"msg_{uuid4().hex}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
}
],
"usage": {
"input_tokens": PROMPT_TOKENS,
"output_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
},
}
return Reply(body=json.dumps(response).encode())
return respond
def _rows_for_key(key: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT litellm_call_id, request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(key.encode()).hexdigest(),),
)
@pytest.mark.timeout(120)
def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once(gateway: Gateway) -> None:
prompt_prefix: Final = f"spend once {uuid4().hex[:8]} "
long_prompt: Final = (prompt_prefix + "lorem ipsum " * 30_000)[:300_000]
upstream_id: Final = f"resp_{uuid4().hex}"
with (
wire_server(_responses_reply(long_prompt, upstream_id)) as wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(
model="openai/responses/gpt-6.1-sol",
api_base=f"{wire.url}/v1",
input_cost_per_token=INPUT_RATE,
output_cost_per_token=OUTPUT_RATE,
)
key: Final = scenario.key(models=[model])
digest: Final = sha256(key.encode()).hexdigest()
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": long_prompt}], "stream": False},
key=key,
)
assert response.status_code == 200, response.text
assert "ok" in response.text, response.text
assert response.json()["id"] == upstream_id, response.text
call_id: Final = response.headers["x-litellm-call-id"]
rows: Final = eventually(
lambda: _rows_for_key(key),
lambda values: any(row["litellm_call_id"] == call_id for row in values),
seconds=70,
)
call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id]
assert len(call_rows) == 1, rows
assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows
logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode()
assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows
key_spend: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= EXPECTED_SPEND,
seconds=70,
)
assert float(str(key_spend[0]["spend"])) == pytest.approx(EXPECTED_SPEND), key_spend

60
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -5811,7 +5811,7 @@ def _interactions_logging_obj(stream: bool, call_type: str = "acreate"):
messages=[],
stream=stream,
call_type=call_type,
start_time=time.time(),
start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc).timestamp(),
litellm_call_id="interactions-call-id",
function_id="interactions-fn-id",
)
@ -9153,3 +9153,61 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch):
logging_module._in_memory_loggers.clear()
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
is_otel_v2_enabled.cache_clear()
@pytest.mark.asyncio
async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch):
"""
The create's in_progress success handler can still be parked on awaited work when
the settlement arrives. The completed result is the only event carrying usage and
cost, so it has to reach the success callbacks whatever the one-success-log-per-request
dedupe does, or the settlement is marked billed with no spend row behind it.
"""
from litellm.litellm_core_utils import litellm_logging
from litellm.types.interactions import InteractionsAPIResponse
class CountingLogger(CustomLogger):
def __init__(self):
super().__init__()
self.logged_results: list[object] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.logged_results.append(response_obj)
counting_logger: Final = CountingLogger()
monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger])
in_progress_parked: Final = asyncio.Event()
release_in_progress: Final = asyncio.Event()
real_truncate: Final = litellm_logging.truncate_base64_in_messages_async
calls: list[int] = []
async def parked_then_real_truncate(messages):
calls.append(1)
if len(calls) == 1:
in_progress_parked.set()
await release_in_progress.wait()
return await real_truncate(messages)
monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", parked_then_real_truncate)
logging_obj: Final = _interactions_logging_obj(stream=False)
in_progress: Final = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress")
completed: Final = InteractionsAPIResponse(
id="interactions/abc",
model="gemini-2.5-flash",
status="completed",
steps=[],
usage=dict(INTERACTIONS_USAGE_BLOCK),
)
first: Final = asyncio.create_task(logging_obj.async_success_handler(result=in_progress))
await in_progress_parked.wait()
completion: Final = asyncio.create_task(logging_obj.async_log_background_interaction_completion(result=completed))
await asyncio.sleep(0)
await asyncio.sleep(0)
release_in_progress.set()
await first
await completion
assert counting_logger.logged_results == [completed, in_progress], counting_logger.logged_results

248
tests/unit/proxy/db/test_model_usage_rollup.py
View file

@ -1,9 +1,68 @@
import asyncio
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
from typing import Any
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.proxy.db.model_usage_rollup import increment_daily_model_usage, model_usage_task_type
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.model_usage_rollup import (
ModelUsageKey,
ModelUsageTransaction,
build_model_usage_transaction,
flush_model_usage_transactions,
model_usage_task_type,
)
class _FakeBatcher:
def __init__(self) -> None:
self.litellm_dailymodelusage = MagicMock()
async def __aenter__(self) -> "_FakeBatcher":
return self
async def __aexit__(self, *args: Any) -> None:
return None
def _prisma(batch_: MagicMock) -> MagicMock:
prisma = MagicMock()
prisma.db.batch_ = batch_
return prisma
def _payload(**overrides: Any) -> dict[str, Any]:
return {
"spend": 0.25,
"prompt_tokens": 10,
"completion_tokens": 20,
"startTime": datetime(2026, 9, 28, 13, tzinfo=timezone.utc),
"model": "openai/gpt-5.4-mini",
"model_group": "fast-chat",
"metadata": "{}",
"request_tags": "[]",
"custom_llm_provider": "openai",
"status": "success",
**overrides,
}
def _transaction(model: str, spend: float, successful: bool = True) -> ModelUsageTransaction:
return ModelUsageTransaction(
key=ModelUsageKey(
date="2026-09-28", model_group=model, model=model, custom_llm_provider="openai", task_type="debugging"
),
spend=spend,
prompt_tokens=10,
completion_tokens=5,
successful=successful,
)
async def _no_sleep(seconds: float) -> None:
return None
def test_model_usage_task_type_reads_task_tag_or_defaults() -> None:
@ -14,76 +73,131 @@ def test_model_usage_task_type_reads_task_tag_or_defaults() -> None:
assert model_usage_task_type("not json") == "uncategorized"
@pytest.mark.asyncio
async def test_increment_daily_model_usage_uses_atomic_prisma_upsert() -> None:
table = MagicMock()
table.upsert = AsyncMock()
prisma_client = MagicMock()
prisma_client.db.litellm_dailymodelusage = table
payload = {
"request_id": "request-1",
"call_type": "acompletion",
"api_key": "key",
"spend": 0.25,
"total_tokens": 30,
"prompt_tokens": 10,
"completion_tokens": 20,
"startTime": datetime(2026, 9, 28, tzinfo=timezone.utc),
"endTime": datetime(2026, 9, 28, tzinfo=timezone.utc),
"completionStartTime": None,
"model": "openai/gpt-5.4-mini",
"model_id": None,
"model_group": "fast-chat",
"mcp_namespaced_tool_name": None,
"agent_id": None,
"api_base": "",
"user": "user",
"metadata": "{}",
"cache_hit": "False",
"cache_key": "",
"request_tags": "[]",
"team_id": None,
"organization_id": None,
"end_user": None,
"requester_ip_address": None,
"custom_llm_provider": "openai",
"messages": None,
"response": None,
"proxy_server_request": None,
"session_id": None,
"request_duration_ms": 20,
"status": "success",
"litellm_call_id": None,
}
def test_build_model_usage_transaction_keys_on_day_model_and_task() -> None:
transaction = build_model_usage_transaction(_payload(request_tags='["task:debugging"]', status="failure"))
await increment_daily_model_usage(prisma_client, payload)
assert transaction == ModelUsageTransaction(
key=ModelUsageKey(
date="2026-09-28",
model_group="fast-chat",
model="openai/gpt-5.4-mini",
custom_llm_provider="openai",
task_type="debugging",
),
spend=0.25,
prompt_tokens=10,
completion_tokens=20,
successful=False,
)
call = table.upsert.await_args.kwargs
assert call["data"]["create"]["request_count"] == 1
assert call["data"]["update"]["completion_tokens"] == {"increment": 20}
assert call["data"]["create"]["task_type"] == "uncategorized"
def test_build_model_usage_transaction_falls_back_for_missing_model_fields() -> None:
transaction = build_model_usage_transaction(
_payload(model="", model_group=None, custom_llm_provider=None, startTime="2026-09-28T01:02:03Z")
)
assert transaction is not None
assert transaction.key == ModelUsageKey(
date="2026-09-28",
model_group="unknown",
model="unknown",
custom_llm_provider="unknown",
task_type="uncategorized",
)
@pytest.mark.parametrize(
"overrides",
[{"metadata": '{"internal_call_origin": "health_check"}'}, {"startTime": "bad"}],
)
def test_build_model_usage_transaction_skips_internal_calls_and_bad_dates(overrides: dict[str, Any]) -> None:
assert build_model_usage_transaction(_payload(**overrides)) is None
@pytest.mark.asyncio
async def test_increment_daily_model_usage_records_task_from_request_tags() -> None:
table = MagicMock()
table.upsert = AsyncMock()
prisma_client = MagicMock()
prisma_client.db.litellm_dailymodelusage = table
payload = {
"call_type": "acompletion",
"spend": 0.1,
"prompt_tokens": 1,
"completion_tokens": 2,
"startTime": datetime(2026, 9, 28, tzinfo=timezone.utc),
"model": "gpt-5",
"model_group": "gpt-5",
"metadata": "{}",
"request_tags": '["task:debugging"]',
"custom_llm_provider": "openai",
"status": "success",
async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None:
batcher = _FakeBatcher()
prisma = _prisma(MagicMock(return_value=batcher))
await flush_model_usage_transactions(
prisma_client=prisma,
transactions=[
_transaction("gpt-5", 0.5),
_transaction("claude", 1.0),
_transaction("gpt-5", 0.25, successful=False),
_transaction("gpt-5", 0.25),
],
)
upserts = {
call.kwargs["where"]["date_model_group_model_custom_llm_provider_task_type"]["model"]: call.kwargs["data"]
for call in batcher.litellm_dailymodelusage.upsert.call_args_list
}
assert list(upserts) == ["claude", "gpt-5"]
gpt = upserts["gpt-5"]
assert gpt["create"]["spend"] == 1.0
assert gpt["create"]["prompt_tokens"] == 30
assert gpt["create"]["completion_tokens"] == 15
assert gpt["create"]["request_count"] == 3
assert gpt["create"]["successful_requests"] == 2
assert gpt["create"]["failed_requests"] == 1
assert gpt["update"] == {
"spend": {"increment": 1.0},
"prompt_tokens": {"increment": 30},
"completion_tokens": {"increment": 15},
"request_count": {"increment": 3},
"successful_requests": {"increment": 2},
"failed_requests": {"increment": 1},
}
assert upserts["claude"]["create"]["request_count"] == 1
await increment_daily_model_usage(prisma_client, payload)
assert table.upsert.await_args.kwargs["data"]["create"]["task_type"] == "debugging"
@pytest.mark.asyncio
async def test_flush_with_no_transactions_touches_nothing() -> None:
prisma = _prisma(MagicMock())
await flush_model_usage_transactions(prisma_client=prisma, transactions=[])
prisma.db.batch_.assert_not_called()
@pytest.mark.asyncio
async def test_flush_retries_connection_errors(monkeypatch: pytest.MonkeyPatch) -> None:
batcher = _FakeBatcher()
prisma = _prisma(MagicMock(side_effect=[httpx.ConnectError("down"), batcher]))
monkeypatch.setattr("litellm.proxy.db.model_usage_rollup.asyncio.sleep", _no_sleep)
await flush_model_usage_transactions(prisma_client=prisma, transactions=[_transaction("gpt-5", 0.1)])
assert prisma.db.batch_.call_count == 2
batcher.litellm_dailymodelusage.upsert.assert_called_once()
@pytest.mark.asyncio
async def test_flush_does_not_retry_ambiguous_errors() -> None:
prisma = _prisma(MagicMock(side_effect=httpx.ReadTimeout("ambiguous")))
with pytest.raises(httpx.ReadTimeout):
await flush_model_usage_transactions(prisma_client=prisma, transactions=[_transaction("gpt-5", 0.1)])
prisma.db.batch_.assert_called_once()
@pytest.mark.asyncio
async def test_request_time_path_queues_usage_instead_of_writing_to_the_db() -> None:
prisma = MagicMock()
prisma.model_usage_transactions = []
prisma._model_usage_transactions_lock = asyncio.Lock()
await DBSpendUpdateWriter()._batch_database_updates(
response_cost=0.25,
user_id="u1",
hashed_token="t1",
team_id=None,
org_id=None,
end_user_id=None,
prisma_client=prisma,
litellm_proxy_budget_name=None,
payload=_payload(request_id="req-1"),
)
assert [transaction.key.model for transaction in prisma.model_usage_transactions] == ["openai/gpt-5.4-mini"]
prisma.db.litellm_dailymodelusage.upsert.assert_not_called()

23
tests/unit/proxy/management_endpoints/test_model_insights_endpoints.py
View file

@ -7,7 +7,7 @@ from fastapi.testclient import TestClient
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.model_usage_rollup import increment_daily_model_usage
from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction, flush_model_usage_transactions
from litellm.proxy.management_endpoints.model_insights_endpoints import router
@ -199,7 +199,7 @@ class _InMemoryUsageTable:
def __init__(self) -> None:
self.rows: dict[tuple[str, ...], dict[str, float]] = {}
async def upsert(self, where: dict, data: dict) -> None:
def upsert(self, where: dict, data: dict) -> None:
key_fields = where["date_model_group_model_custom_llm_provider_task_type"]
key = tuple(key_fields.values())
if key not in self.rows:
@ -221,11 +221,23 @@ class _InMemoryUsageTable:
return list(grouped.values())
class _InMemoryBatcher:
def __init__(self, table: _InMemoryUsageTable) -> None:
self.litellm_dailymodelusage = table
async def __aenter__(self) -> "_InMemoryBatcher":
return self
async def __aexit__(self, *args: object) -> None:
return None
@pytest.mark.asyncio
async def test_model_insights_reads_back_what_the_rollup_wrote() -> None:
table = _InMemoryUsageTable()
prisma = MagicMock()
prisma.db.litellm_dailymodelusage = table
prisma.db.batch_ = MagicMock(return_value=_InMemoryBatcher(table))
payload = {
"call_type": "acompletion",
"spend": 0.5,
@ -240,8 +252,11 @@ async def test_model_insights_reads_back_what_the_rollup_wrote() -> None:
"status": "success",
}
await increment_daily_model_usage(prisma, payload)
await increment_daily_model_usage(prisma, {**payload, "request_tags": "[]"})
transactions = (
build_model_usage_transaction(payload),
build_model_usage_transaction({**payload, "request_tags": "[]"}),
)
await flush_model_usage_transactions(prisma, [t for t in transactions if t is not None])
body = _call(table, "metric=requests").json()

2
tests/unit/proxy/test_update_spend.py
View file

@ -36,6 +36,7 @@ class MockPrismaClient:
self.spend_log_transactions = []
self.daily_user_spend_transactions = {}
self.tool_usage_transactions = []
self.model_usage_transactions = []
self.autorouter_turn_transactions = []
self.baseline_accounting_transactions = []
self.baseline_accounting_lock = asyncio.Lock()
@ -49,6 +50,7 @@ class MockPrismaClient:
self._spend_log_transactions_lock = asyncio.Lock()
self.spend_log_write_lock = asyncio.Lock()
self._tool_usage_transactions_lock = asyncio.Lock()
self._model_usage_transactions_lock = asyncio.Lock()
self._autorouter_turn_transactions_lock = asyncio.Lock()
def jsonify_object(self, obj):

2
tests/unit/proxy/utils/prisma_and_spend/conftest.py
View file

@ -133,6 +133,8 @@ def mock_prisma_client() -> MagicMock:
client.spend_log_write_lock = asyncio.Lock()
client.tool_usage_transactions = []
client._tool_usage_transactions_lock = asyncio.Lock()
client.model_usage_transactions = []
client._model_usage_transactions_lock = asyncio.Lock()
client.jsonify_object = lambda data: dict(data)
client.db.is_connected = MagicMock(return_value=False)
client.db.connect = AsyncMock()

27
tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py
View file

@ -223,6 +223,33 @@ async def test_update_spend_logs_job_drains_tool_queue_when_spend_queue_empty(
assert mock_prisma_client.tool_usage_transactions == []
@pytest.mark.asyncio
async def test_update_spend_logs_job_drains_the_whole_model_usage_queue_in_one_run(
mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
import litellm.proxy.db.model_usage_rollup as model_usage_mod
import litellm.proxy.db.spend_log_tool_index as tool_mod
import litellm.proxy.guardrails.usage_tracking as guard_mod
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
queued = [MagicMock() for _ in range(25_000)]
mock_prisma_client.model_usage_transactions = list(queued)
monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False)
monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False)
flush_stub = AsyncMock()
monkeypatch.setattr(model_usage_mod, "flush_model_usage_transactions", flush_stub, raising=False)
await update_spend_logs_job(
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging,
)
assert flush_stub.await_args.kwargs["transactions"] == queued
assert mock_prisma_client.model_usage_transactions == []
@pytest.mark.asyncio
async def test_update_spend_logs_job_processes_and_clears_queue(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch

64
tests/unit/test_utils.py
View file

@ -6513,3 +6513,67 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None:
assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}]
assert payload not in str(logged)
@pytest.mark.asyncio
async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None:
"""Chat over the Responses bridge exits two @client wrappers with one logging object. Issue
#44500: both exits enqueued a success handler, and on a large prompt the first one yielded to
the worker-thread base64 offload before it marked ``has_logged_async_success``, so the second
passed the check too and the request was logged and billed twice. The schedule step claims the
log for the object synchronously, so only the inner provider-shaped result is ever logged."""
from litellm.litellm_core_utils import litellm_logging
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.utils import _dispatch_success_logging
class CountingLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.logged_results: list[object] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
if kwargs["litellm_call_id"] == "bridge-call-id":
self.logged_results.append(response_obj)
counting_logger: Final = CountingLogger()
monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger])
real_truncate: Final = litellm_logging.truncate_base64_in_messages_async
async def yielding_truncate(messages):
await asyncio.sleep(0)
return await real_truncate(messages)
monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", yielding_truncate)
messages: Final = [{"role": "user", "content": "hello"}]
logging_obj: Final = Logging(
model="gpt-5.6-luna",
messages=messages,
stream=False,
call_type="acompletion",
start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(),
litellm_call_id="bridge-call-id",
function_id="bridge-fn-id",
)
logging_obj.update_environment_variables(
litellm_params={}, optional_params={}, model="gpt-5.6-luna", custom_llm_provider="openai", input=messages
)
inner_result: Final = ModelResponse(id="inner")
outer_result: Final = ModelResponse(id="outer")
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
for result in (inner_result, outer_result):
_dispatch_success_logging(
logging_obj=logging_obj,
result=result,
start_time=now,
end_time=now,
is_completion_with_fallbacks=False,
is_litellm_internal_call=False,
)
await asyncio.sleep(0)
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
assert len(counting_logger.logged_results) == 1, counting_logger.logged_results
assert counting_logger.logged_results[0] is inner_result

2
uv.lock generated
View file

@ -4500,7 +4500,7 @@ wheels = [
[[package]]
name = "litellm"
version = "1.105.0"
version = "1.105.1"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },