diff --git a/db_scripts/backfill_daily_tool_spend.sql b/db_scripts/backfill_daily_tool_spend.sql new file mode 100644 index 00000000000..309b9dbe0ff --- /dev/null +++ b/db_scripts/backfill_daily_tool_spend.sql @@ -0,0 +1,46 @@ +-- One-shot backfill of the LiteLLM_DailyToolSpend rollup from the per-request +-- LiteLLM_SpendLogToolIndex x LiteLLM_SpendLogs tables. +-- +-- This is an opt-in, manual operation. New deployments do not need it: the +-- rollup is written at request time from the moment the release is deployed. +-- Run it only if you want the Cost Optimization "Spend by tool" card to show +-- history from before the deploy, and only once. +-- +-- IMPORTANT caveats before running: +-- +-- 1. Pre-deploy index rows may include tools that were merely DECLARED in a +-- request body but never invoked (the release this ships with stops +-- recording those). For agentic clients that declare many tools per +-- request, backfilled history attributes each request's full spend to +-- every declared tool, overstating per-tool spend. Post-deploy rows do not +-- have this problem. If your traffic is mostly such clients, consider not +-- backfilling. +-- +-- 2. Coverage is bounded by spend-log retention: rows older than +-- maximum_spend_logs_retention_period are already gone. +-- +-- 3. Replace the cutover timestamp below with the time you deployed the +-- release, so backfilled per-request rows cannot double-count on top of +-- rollup rows the new writer already created. ON CONFLICT DO NOTHING is a +-- second guard for (date, tool_name) buckets the writer already touched: +-- such buckets keep the writer's numbers and skip the backfill's. +-- +-- Usage: +-- psql "$DATABASE_URL" -v cutover="'2026-07-25T00:00:00Z'" -f db_scripts/backfill_daily_tool_spend.sql + +SET TIME ZONE 'UTC'; + +INSERT INTO "LiteLLM_DailyToolSpend" (date, tool_name, spend, total_tokens, request_count, created_at, updated_at) +SELECT + to_char(ti.start_time, 'YYYY-MM-DD') AS date, + ti.tool_name, + COALESCE(SUM(sl.spend), 0) AS spend, + COALESCE(SUM(sl.total_tokens), 0) AS total_tokens, + COUNT(*) AS request_count, + now() AS created_at, + now() AS updated_at +FROM "LiteLLM_SpendLogToolIndex" ti +JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id +WHERE ti.start_time < :cutover::timestamptz +GROUP BY 1, 2 +ON CONFLICT (date, tool_name) DO NOTHING; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql new file mode 100644 index 00000000000..e02ed01a554 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql @@ -0,0 +1,12 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyToolSpend" ( + "date" TEXT NOT NULL, + "tool_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "total_tokens" BIGINT NOT NULL DEFAULT 0, + "request_count" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyToolSpend_pkey" PRIMARY KEY ("date","tool_name") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..1014b472c61 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1457,7 +1457,7 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) -TOOL_SPEND_MAX_WINDOW_DAYS = 30 +TOOL_SPEND_TOP_TOOLS = 100 SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index f7ff4d6b16f..4e3d94e2ab3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,6 +6,7 @@ import mimetypes import re import xml.etree.ElementTree as ET from enum import Enum +from collections.abc import Mapping from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -5350,7 +5351,9 @@ def prompt_factory( def get_attribute_or_key(tool_or_function, attribute, default=None): if hasattr(tool_or_function, attribute): return getattr(tool_or_function, attribute) - return tool_or_function.get(attribute, default) + if isinstance(tool_or_function, Mapping): + return tool_or_function.get(attribute, default) + return default class NormalizedToolCall(TypedDict): @@ -5379,14 +5382,18 @@ def _parse_tool_call_arguments(raw: Any, tool_name: Optional[str], context: str) return parsed if isinstance(parsed, dict) else {} -def _tool_calls_from_chat_completion_response(response: Any) -> list[NormalizedToolCall]: +def _tool_calls_from_chat_completion_response( + response: Any, include_all_choices: bool = False +) -> list[NormalizedToolCall]: choices = get_attribute_or_key(response, "choices", None) if not (isinstance(choices, list) and choices): return [] - message = get_attribute_or_key(choices[0], "message", None) - tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None - if not isinstance(tool_calls, list): - return [] + tool_calls: list[Any] = [] + for choice in choices if include_all_choices else choices[:1]: + message = get_attribute_or_key(choice, "message", None) + choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None + if isinstance(choice_tool_calls, list): + tool_calls.extend(choice_tool_calls) result: list[NormalizedToolCall] = [] for tc in tool_calls: fn = get_attribute_or_key(tc, "function", None) @@ -5449,7 +5456,7 @@ def _tool_calls_from_anthropic_messages_response(response: Any) -> list[Normaliz return result -def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: +def get_tool_calls_from_response(response: Any, include_all_choices: bool = False) -> list[NormalizedToolCall]: """ Extract tool/function calls from a response object into a normalized ``{"id", "name", "arguments"}`` shape, regardless of which API surface @@ -5457,11 +5464,20 @@ def get_tool_calls_from_response(response: Any) -> list[NormalizedToolCall]: the Responses API (``output`` items of type ``function_call``), or the Anthropic Messages API (``content`` blocks of type ``tool_use``). + ``include_all_choices`` decides the chat-completions scope: the default + reads only ``choices[0]``, which is what consumers that act on THE reply + (e.g. guardrails rebuilding the primary assistant message) want; usage + accounting passes True because every choice of an ``n>1`` request costs + money and its tool calls really ran. The other surfaces have a single + output, so the flag has no effect on them. + Callers that only care about a specific tool should filter the result by ``name`` themselves -- this returns every tool call found. """ + chat_tool_calls = _tool_calls_from_chat_completion_response(response, include_all_choices=include_all_choices) + if chat_tool_calls: + return chat_tool_calls for extractor in ( - _tool_calls_from_chat_completion_response, _tool_calls_from_responses_api_response, _tool_calls_from_anthropic_messages_response, ): diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 972c831073f..96f6ee89d56 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -26858,12 +26858,6 @@ } ], "title": "Start Date" - }, - "total_spend": { - "default": 0.0, - "description": "Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist", - "title": "Total Spend", - "type": "number" } }, "title": "ToolSpendResponse", @@ -27417,7 +27411,7 @@ }, "/v1/tool/spend": { "get": { - "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nJoins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to\n``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools\ncounts its full spend toward each of those tools, so per-tool numbers are\nattributions. ``total_spend`` is the deduplicated spend of every request that\ncalled at least one tool in the window, so it never double counts.\n\n``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to\n31 calendar dates inclusive, the same width as the endpoint's default window):\na wider requested range is clamped, and the response's ``start_date`` reflects\nthe effective window actually served.", + "description": "Spend attributed to each tool over a date range, for the Cost Optimization dashboard.\n\nReads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked\ntools only (MCP tool calls and response tool_calls; declaring a tool without\ninvoking it does not count). A request that invoked multiple tools counts its\nfull spend toward each of them, so per-tool numbers are attributions and do not\nsum to a deduplicated total.\n\n``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in\nSQL, and ``daily`` covers only those tools, so the response is bounded by\ndays x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many\ndistinct tool names exist.", "operationId": "get_tool_spend_v1_tool_spend_get", "parameters": [ { @@ -27588,7 +27582,7 @@ }, "/v1/tool/{tool_name}/logs": { "get": { - "description": "Return paginated spend logs for requests that used this tool (from SpendLogToolIndex).", + "description": "Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex).\nDeclaring a tool in a request body without the model invoking it does not create an entry.", "operationId": "get_tool_usage_logs_v1_tool__tool_name__logs_get", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 20d0d535b87..dcd2de07e07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3639,6 +3639,13 @@ DB_CONNECTION_ERROR_TYPES = ( httpx.ReadTimeout, ) +# What a NON-IDEMPOTENT write (increment upsert) may retry: only ConnectError +# proves the statements never reached the database. Post-send errors are +# ambiguous; a stalled statement can leave its transaction open on the pooled +# connection, where a retry stacks a second increment set into the same commit. +# Idempotent writes (create_many with skip_duplicates) may retry the full tuple. +DB_RETRY_SAFE_ERROR_TYPES = (httpx.ConnectError,) + class SSOUserDefinedValues(TypedDict): models: List[str] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2262141f426..fd8132fef22 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,7 +34,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, DailyEndUserSpendTransaction, @@ -182,6 +182,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, ) + await self._enqueue_tool_usage_transaction( + payload=payload, + completion_response=completion_response, + prisma_client=prisma_client, + kwargs=kwargs, + ) else: verbose_proxy_logger.debug( "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." @@ -223,6 +229,36 @@ class DBSpendUpdateWriter: end_user_id, ) + async def _enqueue_tool_usage_transaction( + self, + payload: SpendLogsPayload, + completion_response: "litellm.ModelResponse | Any | Exception | None", + prisma_client: "PrismaClient | None", + kwargs: "dict | None" = None, + ) -> None: + try: + if prisma_client is None: + return + from litellm.proxy.db.spend_log_tool_index import ( + build_tool_usage_transaction, + ) + + transaction = build_tool_usage_transaction( + request_id=payload["request_id"], + start_time_iso=str(payload["startTime"]), + mcp_namespaced_tool_name=payload.get("mcp_namespaced_tool_name"), + spend=payload["spend"], + total_tokens=payload["total_tokens"], + completion_response=completion_response, + realtime_tool_calls=(kwargs or {}).get("realtime_tool_calls"), + ) + if transaction is None: + return + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions.append(transaction) + except Exception as e: + verbose_proxy_logger.debug("_enqueue_tool_usage_transaction error (non-blocking): %s", e) + def _enqueue_tool_registry_upsert( self, kwargs: Optional[dict], @@ -299,21 +335,10 @@ class DBSpendUpdateWriter: _enqueue(name) # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- - if completion_response is not None and hasattr(completion_response, "choices"): - for choice in completion_response.choices or []: - message = getattr(choice, "message", None) - if message is None: - continue - tool_calls = getattr(message, "tool_calls", None) - if not tool_calls: - continue - for tc in tool_calls: - fn = getattr(tc, "function", None) - if fn is None: - continue - tool_name = getattr(fn, "name", None) - if tool_name: - _enqueue(tool_name) + from litellm.proxy.db.spend_log_tool_index import response_tool_call_names + + for tool_name in response_tool_call_names(completion_response): + _enqueue(tool_name) except Exception as e: verbose_proxy_logger.debug("_enqueue_tool_registry_upsert error (non-blocking): %s", e) @@ -1096,7 +1121,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1139,7 +1164,7 @@ class DBSpendUpdateWriter: }, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1172,7 +1197,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1219,7 +1244,7 @@ class DBSpendUpdateWriter: ) # Transaction succeeded, break out of retry loop break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1261,7 +1286,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, @@ -1347,7 +1372,7 @@ class DBSpendUpdateWriter: data={"spend": {"increment": response_cost}}, ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, @@ -1644,7 +1669,7 @@ class DBSpendUpdateWriter: break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: _raise_failed_update_spend_exception( e=e, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 80036e235f7..802e893d473 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -1,140 +1,150 @@ """ -Track tool usage for the dashboard: insert into SpendLogToolIndex when spend logs -are written, so "last N requests for tool X" and "how is this tool called in production" -queries are fast. +Tool usage tracking for the dashboard. + +At request time the spend writer builds one ToolUsageTransaction per request that +invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not- +invoked tools are excluded) and queues it on the prisma client. The spend-log flush +job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and +LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads) in a +single transaction, so a failed flush never leaves a partial rollup increment. """ +from __future__ import annotations + +import asyncio +import random +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Dict, List, Set +from itertools import groupby +from typing import TYPE_CHECKING, Any, Sequence -from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import SpendLogToolIndexRepository +from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient -def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: - """Extract tool names from OpenAI-style tool_calls list into out.""" - if not isinstance(tool_calls, list): - return - for tc in tool_calls: - if not isinstance(tc, dict): - continue - fn = tc.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - out.add(name.strip()) +@dataclass(frozen=True, slots=True) +class ToolUsageTransaction: + request_id: str + date: str + start_time: datetime + tool_names: tuple[str, ...] + spend: float + total_tokens: int -def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: - """ - Extract deduplicated tool names from a spend log payload. - Sources: mcp_namespaced_tool_name, response (tool_calls), proxy_server_request (tools). - """ - tool_names: Set[str] = set() +def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: + """Tool names invoked in a completion response, in call order, for any response + surface get_tool_calls_from_response understands (chat completions, Responses + API output items, Anthropic Messages tool_use blocks). Reads every choice of + an ``n>1`` chat response: each choice cost money and its tool calls ran.""" + if completion_response is None or isinstance(completion_response, Exception): + return () + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) - # Top-level MCP tool name (single tool per request for that flow) - mcp_name = payload.get("mcp_namespaced_tool_name") - if mcp_name and isinstance(mcp_name, str) and mcp_name.strip(): - tool_names.add(mcp_name.strip()) - - # Response: OpenAI-style tool_calls[].function.name or choices[0].message.tool_calls - response_raw = payload.get("response") - if response_raw: - response_obj = safe_json_loads(response_raw, default=None) if isinstance(response_raw, str) else response_raw - if isinstance(response_obj, dict): - _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) - choices = response_obj.get("choices") - if isinstance(choices, list) and choices: - msg = choices[0].get("message") if isinstance(choices[0], dict) else None - if isinstance(msg, dict): - _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) - - # Request body: tools[].function.name - request_raw = payload.get("proxy_server_request") - if request_raw: - request_obj = safe_json_loads(request_raw, default=None) if isinstance(request_raw, str) else request_raw - if isinstance(request_obj, dict): - body = request_obj.get("body", request_obj) - if isinstance(body, dict): - request_obj = body - if isinstance(request_obj, dict): - tools = request_obj.get("tools") - if isinstance(tools, list): - for t in tools: - if isinstance(t, dict): - fn = t.get("function") - if isinstance(fn, dict): - name = fn.get("name") - if name and isinstance(name, str) and name.strip(): - tool_names.add(name.strip()) - - return tool_names + return tuple( + stripped + for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) + if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) + ) -async def process_spend_logs_tool_usage( - prisma_client: PrismaClient, - logs_to_process: List[Dict[str, Any]], -) -> None: - """ - After spend logs are written: insert SpendLogToolIndex rows from each payload. - Extracts tool names from mcp_namespaced_tool_name, response tool_calls, and - proxy_server_request tools. - """ - if not logs_to_process: - return - - index_rows: List[Dict[str, Any]] = [] - - for payload in logs_to_process: - request_id = payload.get("request_id") - start_time = payload.get("startTime") - if not request_id or not start_time: - continue - if isinstance(start_time, str): - try: - start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if start_time.tzinfo is None: - start_time = start_time.replace(tzinfo=timezone.utc) - - tool_names = _parse_tool_names_from_payload(payload) - for tool_name in tool_names: - index_rows.append( - { - "request_id": request_id, - "tool_name": tool_name, - "start_time": start_time, - } - ) - - if not index_rows: - return - +def build_tool_usage_transaction( + request_id: str, + start_time_iso: str, + mcp_namespaced_tool_name: str | None, + spend: float, + total_tokens: int, + completion_response: Any, + realtime_tool_calls: Any = None, +) -> ToolUsageTransaction | None: + """None when the request invoked no tools. Realtime sessions carry invoked + tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than + on a response object, so they are normalized through the same owner by + wrapping them in the chat-completion shape. Date derivation must match the + daily spend writer's ``startTime.split("T")[0]`` so rollup rows land in the + same UTC day bucket as LiteLLM_DailyUserSpend.""" + mcp_names = ( + (mcp_namespaced_tool_name.strip(),) if mcp_namespaced_tool_name and mcp_namespaced_tool_name.strip() else () + ) + realtime_names = ( + response_tool_call_names({"choices": [{"message": {"tool_calls": realtime_tool_calls}}]}) + if realtime_tool_calls + else () + ) + tool_names = tuple(dict.fromkeys(mcp_names + response_tool_call_names(completion_response) + realtime_names)) + if not tool_names: + return None try: - index_data = [] - for r in index_rows: - st = r["start_time"] - if isinstance(st, str): - try: - st = datetime.fromisoformat(st.replace("Z", "+00:00")) - except (ValueError, TypeError): - continue - if st.tzinfo is None: - st = st.replace(tzinfo=timezone.utc) - index_data.append( - { - "request_id": r["request_id"], - "tool_name": r["tool_name"], - "start_time": st, - } - ) - if index_data: - await SpendLogToolIndexRepository(prisma_client).table.create_many( - data=index_data, - skip_duplicates=True, - ) - except Exception as e: - verbose_proxy_logger.warning("Tool usage tracking (SpendLogToolIndex) failed (non-fatal): %s", e) + start_time = datetime.fromisoformat(start_time_iso.replace("Z", "+00:00")) + except ValueError: + return None + return ToolUsageTransaction( + request_id=request_id, + date=start_time_iso.split("T")[0], + start_time=start_time if start_time.tzinfo else start_time.replace(tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +async def flush_tool_usage_transactions( + prisma_client: PrismaClient, + transactions: Sequence[ToolUsageTransaction], + n_retry_times: int = 3, +) -> None: + """Write index rows and rollup upserts for a drained queue batch in one + transaction. Retries only ConnectError, the one failure that proves the + statements never reached the database. Post-send failures (Read timeouts + and errors) are ambiguous and are NOT retried: the engine can abandon the + transaction open on the pooled connection, so a retry's statements stack + into the same transaction and one commit applies both increment sets. + Ambiguous failures drop the batch; the caller logs it at error. Callers + must not add their own retry around this function.""" + if not transactions: + return + + index_rows = [ + {"request_id": txn.request_id, "tool_name": tool_name, "start_time": txn.start_time} + for txn in transactions + for tool_name in txn.tool_names + ] + per_tool_day = sorted( + ((txn.date, tool_name, txn.spend, txn.total_tokens) for txn in transactions for tool_name in txn.tool_names), + key=lambda entry: (entry[0], entry[1]), + ) + + for attempt in range(n_retry_times + 1): + try: + async with prisma_client.db.batch_() as batcher: + batcher.litellm_spendlogtoolindex.create_many(data=index_rows, skip_duplicates=True) + for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])): + entries = tuple(grouped) + spend = sum(entry[2] for entry in entries) + total_tokens = sum(entry[3] for entry in entries) + batcher.litellm_dailytoolspend.upsert( + where={"date_tool_name": {"date": date_key, "tool_name": tool_name}}, + data={ + "create": { + "date": date_key, + "tool_name": tool_name, + "spend": spend, + "total_tokens": total_tokens, + "request_count": len(entries), + }, + "update": { + "spend": {"increment": spend}, + "total_tokens": {"increment": total_tokens}, + "request_count": {"increment": len(entries)}, + }, + }, + ) + return + except DB_RETRY_SAFE_ERROR_TYPES: + if attempt >= n_retry_times: + raise + await asyncio.sleep(2**attempt + random.uniform(0, 1)) diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index b6a445ef327..ad0b4f7444c 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -11,21 +11,21 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a import uuid from datetime import datetime, timedelta, timezone -from itertools import groupby from typing import TYPE_CHECKING, Annotated, Any, List, Optional from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient from litellm._logging import verbose_proxy_logger -from litellm.constants import TOOL_SPEND_MAX_WINDOW_DAYS +from litellm.constants import TOOL_SPEND_TOP_TOOLS from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( + DailyToolSpendRepository, SpendLogsRepository, SpendLogToolIndexRepository, ) @@ -142,53 +142,18 @@ def _parse_day_start(value: str | None) -> datetime | None: ) -class _ToolSpendRow(BaseModel): - date: str +class _ToolSpendSums(BaseModel): + spend: float = 0.0 + total_tokens: int = 0 + request_count: int = 0 + + +class _TopToolRow(BaseModel): tool_name: str - call_count: int - spend: float - total_tokens: int + sums: _ToolSpendSums = Field(alias="_sum") -class _RequestTotalRow(BaseModel): - total_spend: float - - -_TOOL_SPEND_ROWS = TypeAdapter(list[_ToolSpendRow]) -_REQUEST_TOTAL_ROWS = TypeAdapter(list[_RequestTotalRow]) - - -def _summarize_tool(name: str, grp: tuple[_ToolSpendRow, ...]) -> ToolSpendEntry: - return ToolSpendEntry( - tool_name=name, - spend=sum(r.spend for r in grp), - call_count=sum(r.call_count for r in grp), - total_tokens=sum(r.total_tokens for r in grp), - ) - - -def _build_tool_spend_response( - rows: list[_ToolSpendRow], - total_spend: float, - start_date: str, - end_date: str, -) -> ToolSpendResponse: - daily = [ - ToolSpendDailyEntry(date=r.date, tool_name=r.tool_name, spend=r.spend, call_count=r.call_count) for r in rows - ] - grouped = groupby(sorted(rows, key=lambda r: r.tool_name), key=lambda r: r.tool_name) - by_tool = sorted( - (_summarize_tool(name, tuple(grp)) for name, grp in grouped), - key=lambda e: e.spend, - reverse=True, - ) - return ToolSpendResponse( - by_tool=by_tool, - daily=daily, - total_spend=total_spend, - start_date=start_date, - end_date=end_date, - ) +_TOP_TOOL_ROWS = TypeAdapter(list[_TopToolRow]) @router.get( @@ -205,16 +170,16 @@ async def get_tool_spend( """ Spend attributed to each tool over a date range, for the Cost Optimization dashboard. - Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to - ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools - counts its full spend toward each of those tools, so per-tool numbers are - attributions. ``total_spend`` is the deduplicated spend of every request that - called at least one tool in the window, so it never double counts. + Reads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked + tools only (MCP tool calls and response tool_calls; declaring a tool without + invoking it does not count). A request that invoked multiple tools counts its + full spend toward each of them, so per-tool numbers are attributions and do not + sum to a deduplicated total. - ``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to - 31 calendar dates inclusive, the same width as the endpoint's default window): - a wider requested range is clamped, and the response's ``start_date`` reflects - the effective window actually served. + ``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in + SQL, and ``daily`` covers only those tools, so the response is bounded by + days x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many + distinct tool names exist. """ from litellm.proxy.proxy_server import prisma_client @@ -230,64 +195,46 @@ async def get_tool_spend( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - now = datetime.now(timezone.utc) - end_day = _parse_day_start(end_date) - # Anchor the floor to a midnight so the clamp compares dates with dates: - # parsed start_dates are midnight-aligned, and a floor carrying now's - # time-of-day would invisibly truncate an explicit start_date to mid-day. - today = now.replace(hour=0, minute=0, second=0, microsecond=0) - window_floor = (end_day or today) - timedelta(days=TOOL_SPEND_MAX_WINDOW_DAYS) - start_dt = _parse_day_start(start_date) or window_floor - if start_dt < window_floor: - start_dt = window_floor - end_exclusive = (end_day + timedelta(days=1)) if end_day else now + end_day = _parse_day_start(end_date) or datetime.now(timezone.utc) + start_day = _parse_day_start(start_date) or end_day - timedelta(days=30) + start_str = start_day.strftime("%Y-%m-%d") + end_str = end_day.strftime("%Y-%m-%d") + date_window = {"date": {"gte": start_str, "lte": end_str}} - # ti.start_time defines the window in both queries; the sl."startTime" bounds - # exist only so the planner can use the SpendLogs startTime index, and carry a - # 1s margin because the two writers can disagree by ~1ms on the same request. - rows = await prisma_client.db.query_raw( - """ - SELECT to_char(ti.start_time, 'YYYY-MM-DD') AS date, - ti.tool_name AS tool_name, - COUNT(*)::int AS call_count, - COALESCE(SUM(sl.spend), 0)::double precision AS spend, - COALESCE(SUM(sl.total_tokens), 0)::bigint AS total_tokens - FROM "LiteLLM_SpendLogToolIndex" ti - JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id - WHERE ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') - AND sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - GROUP BY date, ti.tool_name - ORDER BY date ASC, spend DESC - """, - start_dt.isoformat(), - end_exclusive.isoformat(), - ) - totals = await prisma_client.db.query_raw( - """ - SELECT COALESCE(SUM(sl.spend), 0)::double precision AS total_spend - FROM "LiteLLM_SpendLogs" sl - WHERE sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') - interval '1 second' - AND sl."startTime" < ($2::timestamptz AT TIME ZONE 'UTC') + interval '1 second' - AND EXISTS ( - SELECT 1 - FROM "LiteLLM_SpendLogToolIndex" ti - WHERE ti.request_id = sl.request_id - AND ti.start_time >= ($1::timestamptz AT TIME ZONE 'UTC') - AND ti.start_time < ($2::timestamptz AT TIME ZONE 'UTC') + table = DailyToolSpendRepository(prisma_client).table + top_tools = _TOP_TOOL_ROWS.validate_python( + await table.group_by( + by=["tool_name"], + sum={"spend": True, "total_tokens": True, "request_count": True}, + where=date_window, + order={"_sum": {"spend": "desc"}}, + take=TOOL_SPEND_TOP_TOOLS, ) - """, - start_dt.isoformat(), - end_exclusive.isoformat(), + or [] ) - total_rows = _REQUEST_TOTAL_ROWS.validate_python(totals or []) - return _build_tool_spend_response( - rows=_TOOL_SPEND_ROWS.validate_python(rows or []), - total_spend=total_rows[0].total_spend if total_rows else 0.0, - start_date=start_dt.strftime("%Y-%m-%d"), - end_date=(end_day or now).strftime("%Y-%m-%d"), + by_tool = [ + ToolSpendEntry( + tool_name=row.tool_name, + spend=row.sums.spend, + call_count=row.sums.request_count, + total_tokens=row.sums.total_tokens, + ) + for row in top_tools + ] + + daily_rows = ( + await table.find_many( + where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, + order=[{"date": "asc"}, {"spend": "desc"}], + ) + if top_tools + else [] ) + daily = [ + ToolSpendDailyEntry(date=row.date, tool_name=row.tool_name, spend=row.spend, call_count=row.request_count) + for row in daily_rows + ] + return ToolSpendResponse(by_tool=by_tool, daily=daily, start_date=start_str, end_date=end_str) @router.get( @@ -388,7 +335,8 @@ async def get_tool_usage_logs( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex). + Declaring a tool in a request body without the model invoking it does not create an entry. """ from litellm.proxy.proxy_server import prisma_client diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e85ccf150d2..924189fed4b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -41,6 +41,7 @@ from litellm.constants import ( ) from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -175,6 +176,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction Span = Union[_Span, Any] else: @@ -2917,6 +2919,8 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() + tool_usage_transactions: List["ToolUsageTransaction"] = [] + _tool_usage_transactions_lock = asyncio.Lock() def __init__( self, @@ -5334,7 +5338,7 @@ class ProxyUpdateSpend: ) break - except DB_CONNECTION_ERROR_TYPES as e: + except DB_RETRY_SAFE_ERROR_TYPES as e: if i >= n_retry_times: # If we've reached the maximum number of retries _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj @@ -5473,12 +5477,15 @@ async def update_spend( queue_size = len(prisma_client.spend_log_transactions) verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) + async with prisma_client._tool_usage_transactions_lock: + tool_usage_queue_size = len(prisma_client.tool_usage_transactions) + # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. # See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior. # Safe to keep: under high concurrency this can take up to ~30s to run, # so it's unlikely to overlap with monitor_spend_logs_queue. - if queue_size > 0: + if queue_size > 0 or tool_usage_queue_size > 0: await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -5545,10 +5552,14 @@ async def update_spend_logs_job( n_retry_times = 3 MAX_LOGS_PER_INTERVAL = 10000 - # Atomically pop batch from queue + # Atomically pop batch from queue. The tool usage queue counts toward the + # emptiness check: a spend-log write failure aborts a run before the tool + # drain below, and those entries must not strand once the spend queue drains. async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - if queue_size == 0: + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + if queue_size == 0 and tool_queue_size == 0: return async with prisma_client._spend_log_transactions_lock: @@ -5579,17 +5590,23 @@ async def update_spend_logs_job( guardrail_tracking_err, ) - # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X" + # Tool usage tracking: drain the request-time queue into the tool index and the + # LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently + # absent from the rollup, so failures log at error. + async with prisma_client._tool_usage_transactions_lock: + tool_usage_to_process = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL] + prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :] try: - from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage + from litellm.proxy.db.spend_log_tool_index import flush_tool_usage_transactions - await process_spend_logs_tool_usage( + await flush_tool_usage_transactions( prisma_client=prisma_client, - logs_to_process=logs_to_process, + transactions=tool_usage_to_process, ) except Exception as tool_tracking_err: - verbose_proxy_logger.warning( - "Spend tracking - tool usage tracking failed (non-fatal): %s", + verbose_proxy_logger.error( + "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", + len(tool_usage_to_process), tool_tracking_err, ) @@ -5625,9 +5642,13 @@ async def _monitor_spend_logs_queue( while True: try: - # Check queue size with lock protection + # Check queue sizes with lock protection; the tool usage queue keeps + # the monitor firing when a prior failed run left it nonempty. async with prisma_client._spend_log_transactions_lock: - queue_size = len(prisma_client.spend_log_transactions) + spend_queue_size = len(prisma_client.spend_log_transactions) + async with prisma_client._tool_usage_transactions_lock: + tool_queue_size = len(prisma_client.tool_usage_transactions) + queue_size = spend_queue_size + tool_queue_size if queue_size > 0: if queue_size >= threshold: diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 4451f0865da..29c953e06cf 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -23,6 +23,7 @@ from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, DailyPolicyMetricsRepository, DailyTagSpendRepository, + DailyToolSpendRepository, DeletedTeamRepository, DeletedVerificationTokenRepository, DeprecatedVerificationTokenRepository, @@ -104,6 +105,7 @@ __all__ = [ "ManagedVectorStoreIndexRepository", "WorkflowMessageRepository", "DailyTagSpendRepository", + "DailyToolSpendRepository", "SpendLogToolIndexRepository", "SpendLogGuardrailIndexRepository", "UserNotificationsRepository", diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index dc2a7d25259..54008c0950c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -181,6 +181,10 @@ class SpendLogToolIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogtoolindex" +class DailyToolSpendRepository(PrismaTableRepository): + table_name = "litellm_dailytoolspend" + + class SpendLogGuardrailIndexRepository(PrismaTableRepository): table_name = "litellm_spendlogguardrailindex" diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py index 71ec412e8ef..ccf4b7dbc9f 100644 --- a/litellm/types/tool_management.py +++ b/litellm/types/tool_management.py @@ -124,12 +124,5 @@ class ToolSpendDailyEntry(BaseModel): class ToolSpendResponse(BaseModel): by_tool: List[ToolSpendEntry] = Field(default_factory=list) daily: List[ToolSpendDailyEntry] = Field(default_factory=list) - total_spend: float = Field( - 0.0, - description=( - "Deduplicated spend of every request that called at least one tool in the window; " - "less than the sum of per-tool attributed spend whenever multi-tool requests exist" - ), - ) start_date: str | None = None end_date: str | None = None diff --git a/schema.prisma b/schema.prisma index 6713b212314..37ea55f8c13 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1097,6 +1097,19 @@ model LiteLLM_SpendLogToolIndex { @@index([start_time]) } +// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs +model LiteLLM_DailyToolSpend { + date String + tool_name String + spend Float @default(0.0) + total_tokens BigInt @default(0) + request_count BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([date, tool_name]) +} + // Prompt table for storing prompt configurations model LiteLLM_PromptTable { id String @id @default(uuid()) diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index e2dca0a0f81..131f46a3e21 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -28,11 +28,13 @@ class MockPrismaClient: # Initialize transaction lists self.spend_log_transactions = [] self.daily_user_spend_transactions = {} + self.tool_usage_transactions = [] - # Add lock for spend_log_transactions (matches real PrismaClient) + # Add locks for the transaction queues (matches real PrismaClient) import asyncio self._spend_log_transactions_lock = asyncio.Lock() + self._tool_usage_transactions_lock = asyncio.Lock() def jsonify_object(self, obj): return obj diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bcda88ea609..9565de1139c 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3166,3 +3166,34 @@ async def test_bedrock_converse_message_level_cache_point_preserves_ttl_async(): ) assert _collect_cache_points(result) == [{"type": "default", "ttl": "1h"}] + + +def _n_choices_response(*names_per_choice): + from types import SimpleNamespace + + choices = [ + SimpleNamespace( + message=SimpleNamespace( + tool_calls=[SimpleNamespace(id=f"c{i}", function=SimpleNamespace(name=name, arguments="{}"))] + ) + ) + for i, name in enumerate(names_per_choice) + ] + return SimpleNamespace(choices=choices) + + +def test_get_tool_calls_from_response_defaults_to_primary_choice_only(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + assert [tc["name"] for tc in get_tool_calls_from_response(response)] == ["tool_alpha"] + + +def test_get_tool_calls_from_response_include_all_choices_reads_every_choice(): + from litellm.litellm_core_utils.prompt_templates.factory import get_tool_calls_from_response + + response = _n_choices_response("tool_alpha", "tool_beta") + + names = [tc["name"] for tc in get_tool_calls_from_response(response, include_all_choices=True)] + assert names == ["tool_alpha", "tool_beta"] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 8149cf90e70..191080e3a48 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -76,6 +76,162 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert call_args["payload"]["custom_llm_provider"] == "openai" +def _tool_call_response(*names: str) -> object: + from types import SimpleNamespace + + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +def _tool_usage_prisma() -> MagicMock: + prisma = MagicMock() + prisma.tool_usage_transactions = [] + prisma._tool_usage_transactions_lock = asyncio.Lock() + prisma.spend_log_transactions = [] + prisma._spend_log_transactions_lock = asyncio.Lock() + return prisma + + +def _minimal_spend_payload() -> dict: + return { + "request_id": "req-tool-1", + "startTime": datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + "endTime": datetime(2026, 7, 25, 10, 0, 1, tzinfo=timezone.utc), + "spend": 0.0, + "total_tokens": 42, + "mcp_namespaced_tool_name": None, + } + + +@pytest.mark.asyncio +async def test_update_database_enqueues_tool_usage_for_invoked_tools(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + transaction = prisma.tool_usage_transactions[0] + assert transaction.request_id == "req-tool-1" + assert transaction.tool_names == ("get_weather",) + assert transaction.spend == 0.1 + assert transaction.total_tokens == 42 + assert transaction.date == "2026-07-25" + + +@pytest.mark.asyncio +async def test_update_database_enqueues_realtime_tool_usage(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={ + "model": "gpt-realtime", + "realtime_tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "rt_tool", "arguments": "{}"}} + ], + }, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.2, + ) + await asyncio.sleep(0) + + assert len(prisma.tool_usage_transactions) == 1 + assert prisma.tool_usage_transactions[0].tool_names == ("rt_tool",) + + +def test_enqueue_tool_registry_upsert_reads_every_choice(): + from types import SimpleNamespace as NS + + db_writer = DBSpendUpdateWriter() + db_writer.tool_discovery_queue = MagicMock() + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + + db_writer._enqueue_tool_registry_upsert(kwargs={}, completion_response=response) + + enqueued = [call.args[0]["tool_name"] for call in db_writer.tool_discovery_queue.add_update.call_args_list] + assert enqueued == ["tool_alpha", "tool_beta"] + + +@pytest.mark.asyncio +async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): + db_writer = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + db_writer._batch_database_updates = AsyncMock() + prisma = _tool_usage_prisma() + + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=_minimal_spend_payload(), + ), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-4"}, + completion_response=_tool_call_response("get_weather"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.1, + ) + await asyncio.sleep(0) + + assert prisma.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_update_daily_spend_with_null_entity_id(): """ @@ -152,6 +308,84 @@ async def test_update_daily_spend_with_null_entity_id(): assert create_data["failed_requests"] == 0 +def _daily_txn(user_id: str = "user1") -> dict: + return { + "user_id": user_id, + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + +@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 + # sent and its outcome is unknown; the engine can leave the transaction open + # on the pooled connection, so retrying stacks a second set of increments + # into it and one commit applies both. Post-send failures must drop the + # batch (loudly), never retry it. + import httpx + + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=httpx.ReadTimeout("ambiguous")) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + mock_prisma_client.db.batch_.assert_called_once() + + +@pytest.mark.asyncio +async def test_update_daily_spend_retries_connect_errors(monkeypatch): + # ConnectError proves the statements never reached the database, so it is + # the one failure the writer may retry. + import httpx + + mock_batcher = MagicMock() + good_ctx = MagicMock() + good_ctx.__aenter__ = AsyncMock(return_value=mock_batcher) + good_ctx.__aexit__ = AsyncMock(return_value=None) + mock_prisma_client = MagicMock() + mock_prisma_client.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), good_ctx]) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions={"k1": _daily_txn()}, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + assert mock_prisma_client.db.batch_.call_count == 2 + + @pytest.mark.asyncio async def test_update_daily_spend_sorting(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py new file mode 100644 index 00000000000..71073fd216e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -0,0 +1,348 @@ +""" +Tests for the tool usage writer: ToolUsageTransaction construction (invoked tools +only) and the flush that writes LiteLLM_SpendLogToolIndex plus the +LiteLLM_DailyToolSpend rollup in one transaction. +""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.spend_log_tool_index import ( + ToolUsageTransaction, + build_tool_usage_transaction, + flush_tool_usage_transactions, + response_tool_call_names, +) + + +def _response_with_tool_calls(*names: str) -> SimpleNamespace: + tool_calls = [SimpleNamespace(function=SimpleNamespace(name=name)) for name in names] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=tool_calls))]) + + +class _FakeBatcher: + def __init__(self) -> None: + self.litellm_spendlogtoolindex = MagicMock() + self.litellm_dailytoolspend = MagicMock() + + async def __aenter__(self) -> "_FakeBatcher": + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + +def _prisma_with_batcher() -> tuple[MagicMock, _FakeBatcher]: + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(return_value=batcher) + return prisma, batcher + + +class TestBuildToolUsageTransaction: + def test_declared_tools_never_reach_the_transaction(self): + # Regression for the inflation bug: the builder's only non-MCP source is + # the response's tool_calls, so a request declaring N tools while the + # model invokes one produces exactly one attribution. + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_no_invoked_tools_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(tool_calls=None))]), + ) + is None + ) + + def test_mcp_name_and_response_names_dedupe(self): + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("srv/tool_a", "tool_b", "tool_b"), + ) + assert transaction is not None + assert transaction.tool_names == ("srv/tool_a", "tool_b") + + def test_date_matches_daily_spend_writer_derivation(self): + # The daily spend writer derives its date bucket as + # payload["startTime"].split("T")[0] (db_spend_update_writer.py), i.e. the + # timestamp's own calendar date, NOT the astimezone-UTC date. A non-UTC + # isoformat pins the difference: 2026-07-25T22:00:00-07:00 is 2026-07-26 + # in UTC but must bucket as 2026-07-25 to match LiteLLM_DailyUserSpend. + start_time_iso = "2026-07-25T22:00:00-07:00" + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso=start_time_iso, + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + assert transaction is not None + assert transaction.date == start_time_iso.split("T")[0] == "2026-07-25" + + def test_realtime_tool_calls_reach_the_transaction(self): + # Realtime sessions carry invoked tools in kwargs["realtime_tool_calls"] + # (OpenAI tool_calls dict shape, built in realtime_streaming.py), not on a + # response object; they must land in the rollup like any other invocation. + realtime_tool_calls = [ + {"id": "call_1", "type": "function", "function": {"name": "rt_get_weather", "arguments": "{}"}}, + ] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=None, + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("rt_get_weather",) + + def test_realtime_names_dedupe_against_response_names(self): + realtime_tool_calls = [{"type": "function", "function": {"name": "get_weather", "arguments": "{}"}}] + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=_response_with_tool_calls("get_weather"), + realtime_tool_calls=realtime_tool_calls, + ) + assert transaction is not None + assert transaction.tool_names == ("get_weather",) + + def test_n_greater_than_one_tools_from_every_choice_reach_the_transaction(self): + # Regression: an n>1 request pays for every choice, and a tool invoked + # only in a later choice really ran; it must not be dropped because the + # extractor read choices[0] alone. + from types import SimpleNamespace as NS + + response = NS( + choices=[ + NS(message=NS(tool_calls=[NS(function=NS(name="tool_alpha"))])), + NS(message=NS(tool_calls=[NS(function=NS(name="tool_beta"))])), + ] + ) + transaction = build_tool_usage_transaction( + request_id="r1", + start_time_iso="2026-07-25T10:00:00+00:00", + mcp_namespaced_tool_name=None, + spend=0.5, + total_tokens=100, + completion_response=response, + ) + assert transaction is not None + assert transaction.tool_names == ("tool_alpha", "tool_beta") + + def test_unparseable_start_time_returns_none(self): + assert ( + build_tool_usage_transaction( + request_id="r1", + start_time_iso="not-a-timestamp", + mcp_namespaced_tool_name="srv/tool_a", + spend=0.5, + total_tokens=100, + completion_response=None, + ) + is None + ) + + +class TestResponseToolCallNames: + def test_unrecognized_shapes_yield_nothing(self): + assert response_tool_call_names(None) == () + assert response_tool_call_names(SimpleNamespace()) == () + assert response_tool_call_names(ValueError("boom")) == () + + def test_blank_names_are_dropped(self): + assert response_tool_call_names(_response_with_tool_calls(" ", "real_tool")) == ("real_tool",) + + def test_responses_api_output_function_calls(self): + # Regression: /v1/responses carries invocations in output[] items of + # type function_call, not in choices; they must reach the rollup. + response = SimpleNamespace( + output=[ + SimpleNamespace(type="function_call", name="get_weather", call_id="c1", arguments="{}"), + SimpleNamespace(type="message", name=None, call_id=None, arguments=None), + ] + ) + assert response_tool_call_names(response) == ("get_weather",) + + def test_anthropic_messages_tool_use_blocks(self): + response = { + "content": [ + {"type": "text", "text": "checking"}, + {"type": "tool_use", "id": "t1", "name": "ant_get_weather", "input": {"city": "Paris"}}, + ] + } + assert response_tool_call_names(response) == ("ant_get_weather",) + + +def _transaction( + request_id: str, + date: str = "2026-07-25", + tool_names: tuple = ("tool_a",), + spend: float = 1.0, + total_tokens: int = 10, +) -> ToolUsageTransaction: + from datetime import datetime, timezone + + return ToolUsageTransaction( + request_id=request_id, + date=date, + start_time=datetime(2026, 7, 25, 10, 0, tzinfo=timezone.utc), + tool_names=tool_names, + spend=spend, + total_tokens=total_tokens, + ) + + +class TestFlushToolUsageTransactions: + @pytest.mark.asyncio + async def test_multi_tool_request_attributes_full_spend_to_each_tool(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1", tool_names=("tool_a", "tool_b"), spend=0.10, total_tokens=100)], + ) + index_rows = batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["data"] + assert [(r["request_id"], r["tool_name"]) for r in index_rows] == [("r1", "tool_a"), ("r1", "tool_b")] + assert batcher.litellm_spendlogtoolindex.create_many.call_args.kwargs["skip_duplicates"] is True + + upserts = { + c.kwargs["where"]["date_tool_name"]["tool_name"]: c.kwargs["data"] + for c in batcher.litellm_dailytoolspend.upsert.call_args_list + } + assert set(upserts) == {"tool_a", "tool_b"} + for data in upserts.values(): + assert data["create"]["spend"] == 0.10 + assert data["create"]["request_count"] == 1 + assert data["update"]["spend"] == {"increment": 0.10} + assert data["update"]["request_count"] == {"increment": 1} + + @pytest.mark.asyncio + async def test_same_day_same_tool_aggregates_within_batch(self): + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[ + _transaction("r1", spend=0.10, total_tokens=100), + _transaction("r2", spend=0.30, total_tokens=200), + ], + ) + assert batcher.litellm_dailytoolspend.upsert.call_count == 1 + data = batcher.litellm_dailytoolspend.upsert.call_args.kwargs["data"] + assert data["create"] == { + "date": "2026-07-25", + "tool_name": "tool_a", + "spend": pytest.approx(0.40), + "total_tokens": 300, + "request_count": 2, + } + assert data["update"]["spend"] == {"increment": pytest.approx(0.40)} + assert data["update"]["total_tokens"] == {"increment": 300} + assert data["update"]["request_count"] == {"increment": 2} + + @pytest.mark.asyncio + async def test_index_rows_and_rollup_share_one_transaction(self): + # Both writes go through the same batch_() so a failed flush cannot leave + # index rows without their rollup increments (or vice versa); increments + # are not idempotent, so partial states must be unreachable. + prisma, batcher = _prisma_with_batcher() + await flush_tool_usage_transactions( + prisma_client=prisma, + transactions=[_transaction("r1")], + ) + prisma.db.batch_.assert_called_once() + batcher.litellm_spendlogtoolindex.create_many.assert_called_once() + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_empty_batch_touches_nothing(self): + prisma, _ = _prisma_with_batcher() + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[]) + prisma.db.batch_.assert_not_called() + + @pytest.mark.asyncio + async def test_connection_errors_retry_and_succeed(self, monkeypatch): + # A failed batch commits nothing, so retrying a connection error cannot + # double-count; the flush must retry rather than drop the batch. + import httpx + + batcher = _FakeBatcher() + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), batcher]) + sleeps: list[float] = [] + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + assert prisma.db.batch_.call_count == 2 + assert len(sleeps) == 1 + batcher.litellm_dailytoolspend.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_connection_errors_exhaust_retries_then_raise(self, monkeypatch): + import httpx + + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=httpx.ConnectError("down")) + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.spend_log_tool_index.asyncio.sleep", fake_sleep) + with pytest.raises(httpx.ConnectError): + await flush_tool_usage_transactions( + prisma_client=prisma, transactions=[_transaction("r1")], n_retry_times=2 + ) + assert prisma.db.batch_.call_count == 3 + + @pytest.mark.asyncio + async def test_non_connection_errors_do_not_retry(self): + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) + with pytest.raises(ValueError): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("ambiguous_error", ["ReadTimeout", "ReadError"]) + async def test_post_send_ambiguous_errors_drop_without_retry(self, ambiguous_error): + # A ReadTimeout means the statements were sent and the outcome is + # unknown; the engine can leave the transaction open on the pooled + # connection, so a retry's statements would stack into it and one + # commit would apply both increment sets. These must never retry. + import httpx + + error = getattr(httpx, ambiguous_error)("ambiguous") + prisma = MagicMock() + prisma.db.batch_ = MagicMock(side_effect=error) + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) + prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py index c908250fa64..45c3c6c2466 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -19,11 +19,7 @@ from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../..")) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - _build_tool_spend_response, - _ToolSpendRow, - router, -) +from litellm.proxy.management_endpoints.tool_management_endpoints import router from litellm.types.tool_management import LiteLLM_ToolTableRow # --- helpers --- @@ -64,6 +60,30 @@ def _override_auth(): _MOCK_PRISMA = MagicMock() +def _rollup_row(date: str, tool_name: str, spend: float, request_count: int, total_tokens: int) -> MagicMock: + row = MagicMock() + row.date = date + row.tool_name = tool_name + row.spend = spend + row.request_count = request_count + row.total_tokens = total_tokens + return row + + +def _group_row(tool_name: str, spend: float, request_count: int, total_tokens: int) -> dict: + return {"tool_name": tool_name, "_sum": {"spend": spend, "total_tokens": total_tokens, "request_count": request_count}} + + +def _rollup_prisma(group_rows: list, daily_rows: list | None = None) -> MagicMock: + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_spendlogtoolindex.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_dailytoolspend.group_by = AsyncMock(return_value=group_rows) + prisma.db.litellm_dailytoolspend.find_many = AsyncMock(return_value=daily_rows or []) + return prisma + + # --- test class --- @@ -154,21 +174,23 @@ class TestToolManagementEndpoints: assert resp.status_code == 422 def test_tool_spend_route_not_shadowed_by_get_tool(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 assert resp.json()["by_tool"] == [] - def test_tool_spend_aggregates_and_sorts(self): - rows = [ - {"date": "2026-07-01", "tool_name": "search", "call_count": 2, "spend": 1.0, "total_tokens": 100}, - {"date": "2026-07-02", "tool_name": "search", "call_count": 1, "spend": 4.0, "total_tokens": 50}, - {"date": "2026-07-01", "tool_name": "read_file", "call_count": 3, "spend": 2.0, "total_tokens": 300}, + def test_tool_spend_serves_sql_aggregates_and_daily_series(self): + group_rows = [ + _group_row("search", spend=5.0, request_count=3, total_tokens=150), + _group_row("read_file", spend=2.0, request_count=3, total_tokens=300), ] - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(side_effect=[rows, [{"total_spend": 5.5}]]) + daily_rows = [ + _rollup_row("2026-07-01", "search", spend=1.0, request_count=2, total_tokens=100), + _rollup_row("2026-07-01", "read_file", spend=2.0, request_count=3, total_tokens=300), + _rollup_row("2026-07-02", "search", spend=4.0, request_count=1, total_tokens=50), + ] + prisma = _rollup_prisma(group_rows, daily_rows) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 @@ -179,94 +201,89 @@ class TestToolManagementEndpoints: assert search["call_count"] == 3 assert search["total_tokens"] == 150 assert len(body["daily"]) == 3 + assert body["daily"][0]["call_count"] == 2 assert body["start_date"] == "2026-07-01" assert body["end_date"] == "2026-07-02" - assert body["total_spend"] == 5.5 + + def test_tool_spend_coerces_bigint_string_sums(self): + # prisma group_by returns BigInt sums as strings ("808"); the response + # must coerce them to ints rather than 500 on validation. + group_rows = [{"tool_name": "search", "_sum": {"spend": 0.5, "total_tokens": "808", "request_count": "3"}}] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + assert resp.json()["by_tool"][0]["total_tokens"] == 808 + assert resp.json()["by_tool"][0]["call_count"] == 3 + + def test_tool_spend_daily_restricted_to_top_tools_and_capped(self): + from litellm.constants import TOOL_SPEND_TOP_TOOLS + + group_rows = [_group_row("search", spend=5.0, request_count=1, total_tokens=10)] + prisma = _rollup_prisma(group_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + group_kwargs = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs + assert group_kwargs["take"] == TOOL_SPEND_TOP_TOOLS + assert group_kwargs["order"] == {"_sum": {"spend": "desc"}} + daily_where = prisma.db.litellm_dailytoolspend.find_many.await_args.kwargs["where"] + assert daily_where["tool_name"] == {"in": ["search"]} + + def test_tool_spend_skips_daily_query_when_no_tools(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + prisma.db.litellm_dailytoolspend.find_many.assert_not_awaited() @patch("litellm.proxy.proxy_server.prisma_client", None) def test_tool_spend_no_db_returns_500(self): resp = self.client.get("/v1/tool/spend") assert resp.status_code == 500 - def test_tool_spend_end_date_is_inclusive_via_exclusive_next_day_bound(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_reads_rollup_only_never_spendlogs(self): + # Regression for the GA blocker: the dashboard aggregate must be served + # entirely from LiteLLM_DailyToolSpend; any query_raw or SpendLogs table + # access on this path reintroduces the per-request scan. + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 7, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 3, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds + prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_spendlogs.find_many.assert_not_awaited() + prisma.db.litellm_spendlogtoolindex.find_many.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_awaited_once() + + def test_tool_spend_windows_rollup_by_inclusive_date_strings(self): + prisma = _rollup_prisma([]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") + assert resp.status_code == 200 + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-07-01", "lte": "2026-07-02"}} assert resp.json()["end_date"] == "2026-07-02" - def test_tool_spend_start_clamped_to_30_days_before_end(self): - # Clamped floor is end_date minus 30 days, serving up to 31 calendar dates - # inclusive: deliberately the same width as the endpoint's default window, - # so the dashboard's default range never triggers the clamp. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_wide_range_served_fully(self): + # Regression: the 30-day clamp is gone; a 182-day request is served as + # requested because the rollup read is O(tools x dates). + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get("/v1/tool/spend?start_date=2026-01-01&end_date=2026-07-01") assert resp.status_code == 200 - expected_binds = ( - datetime(2026, 6, 1, tzinfo=timezone.utc).isoformat(), - datetime(2026, 7, 2, tzinfo=timezone.utc).isoformat(), - ) - assert prisma.db.query_raw.await_count == 2 - for call in prisma.db.query_raw.await_args_list: - assert tuple(call.args[1:]) == expected_binds - assert resp.json()["start_date"] == "2026-06-01" + where = prisma.db.litellm_dailytoolspend.group_by.await_args.kwargs["where"] + assert where == {"date": {"gte": "2026-01-01", "lte": "2026-07-01"}} + assert resp.json()["start_date"] == "2026-01-01" assert resp.json()["end_date"] == "2026-07-01" - def test_tool_spend_range_within_cap_is_not_clamped(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + def test_tool_spend_defaults_to_trailing_30_days(self): + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-06-25&end_date=2026-07-01") + resp = self.client.get("/v1/tool/spend") assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == datetime(2026, 6, 25, tzinfo=timezone.utc).isoformat() - assert resp.json()["start_date"] == "2026-06-25" - - def test_tool_spend_start_honored_when_end_date_omitted(self): - # Regression: with end_date omitted the floor anchors to today's UTC - # midnight, not now's time-of-day, so an explicit start_date exactly 30 - # days back is served from midnight rather than truncated to mid-day. - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get(f"/v1/tool/spend?start_date={floor_day.strftime('%Y-%m-%d')}") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_clamp_without_end_date_lands_on_midnight(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - floor_day = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=30) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2020-01-01") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - assert call.args[1] == floor_day.isoformat() - assert resp.json()["start_date"] == floor_day.strftime("%Y-%m-%d") - - def test_tool_spend_total_query_bounds_outer_spendlogs_scan(self): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - resp = self.client.get("/v1/tool/spend?start_date=2026-07-01&end_date=2026-07-02") - assert resp.status_code == 200 - for call in prisma.db.query_raw.await_args_list: - sql = call.args[0] - assert 'sl."startTime" >=' in sql - assert 'sl."startTime" <' in sql + today = datetime.now(timezone.utc) + assert resp.json()["end_date"] == today.strftime("%Y-%m-%d") + assert resp.json()["start_date"] == (today - timedelta(days=30)).strftime("%Y-%m-%d") @pytest.mark.parametrize( "query", @@ -279,13 +296,12 @@ class TestToolManagementEndpoints: ], ) def test_tool_spend_malformed_date_returns_400(self, query: str): - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = self.client.get(f"/v1/tool/spend?{query}") assert resp.status_code == 400 assert "Invalid date format" in resp.json()["detail"] - prisma.db.query_raw.assert_not_awaited() + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() def test_tool_spend_non_admin_returns_403(self): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -296,38 +312,8 @@ class TestToolManagementEndpoints: api_key="sk-user", user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER ) client = TestClient(app, raise_server_exceptions=True) - prisma = MagicMock() - prisma.db.query_raw = AsyncMock(return_value=[]) + prisma = _rollup_prisma([]) with patch("litellm.proxy.proxy_server.prisma_client", prisma): resp = client.get("/v1/tool/spend") assert resp.status_code == 403 - prisma.db.query_raw.assert_not_awaited() - - -def _spend_row(date: str, tool_name: str, spend: float, call_count: int = 1, total_tokens: int = 10) -> _ToolSpendRow: - return _ToolSpendRow(date=date, tool_name=tool_name, call_count=call_count, spend=spend, total_tokens=total_tokens) - - -class TestBuildToolSpendResponse: - def test_multi_tool_attribution_double_counts_per_tool_but_not_total(self): - rows = [ - _spend_row("2026-07-01", "a", spend=3.0), - _spend_row("2026-07-01", "b", spend=3.0), - ] - resp = _build_tool_spend_response(rows, total_spend=3.0, start_date="2026-07-01", end_date="2026-07-01") - by_tool = {t.tool_name: t.spend for t in resp.by_tool} - assert by_tool == {"a": 3.0, "b": 3.0} - assert resp.total_spend == 3.0 - - def test_groups_across_days_and_sorts_by_spend(self): - rows = [ - _spend_row("2026-07-01", "b", spend=1.0, call_count=2, total_tokens=100), - _spend_row("2026-07-02", "b", spend=4.0, call_count=1, total_tokens=50), - _spend_row("2026-07-01", "a", spend=2.0, call_count=3, total_tokens=300), - ] - resp = _build_tool_spend_response(rows, total_spend=7.0, start_date="2026-07-01", end_date="2026-07-02") - assert [(t.tool_name, t.spend, t.call_count, t.total_tokens) for t in resp.by_tool] == [ - ("b", 5.0, 3, 150), - ("a", 2.0, 3, 300), - ] - assert len(resp.daily) == 3 + prisma.db.litellm_dailytoolspend.group_by.assert_not_awaited() diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index f969b040a0d..2ba9257e1da 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -193,6 +193,12 @@ async def test_cleanup_old_spend_logs_batch_deletion(): tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0] assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql + # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is + # the only copy of tool spend history once its per-request sources expire, + # so spend-log cleanup must never touch it. + for call in mock_db.execute_raw.call_args_list: + assert "LiteLLM_DailyToolSpend" not in call[0][0] + @pytest.mark.asyncio async def test_cleanup_old_spend_logs_retention_period_cutoff(): diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index af62b7eef62..74c9abd9978 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -128,6 +128,8 @@ def mock_prisma_client() -> MagicMock: client.proxy_logging_obj.failure_handler = AsyncMock() client.spend_log_transactions = [] client._spend_log_transactions_lock = asyncio.Lock() + client.tool_usage_transactions = [] + client._tool_usage_transactions_lock = asyncio.Lock() client.jsonify_object = lambda data: dict(data) client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d5d4de7f2cf..f075acc7307 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -68,12 +68,12 @@ async def test_update_end_user_spend_upserts_each_end_user( @pytest.mark.asyncio -async def test_update_end_user_spend_retries_on_connection_error( +async def test_update_end_user_spend_retries_on_connect_error( mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: - """``DB_CONNECTION_ERROR_TYPES`` failures should be retried with backoff; - once retries are exhausted, ``_raise_failed_update_spend_exception`` is - invoked and the original exception bubbles up. + """``DB_RETRY_SAFE_ERROR_TYPES`` (ConnectError, statements provably never + sent) retries with backoff; once retries are exhausted the original + exception bubbles up via ``_raise_failed_update_spend_exception``. """ import httpx import litellm.proxy.utils as utils_mod @@ -85,11 +85,11 @@ async def test_update_end_user_spend_retries_on_connection_error( monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) - err = httpx.ReadError("conn reset") + err = httpx.ConnectError("down") mock_prisma_client.db.tx = MagicMock(side_effect=err) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() - with pytest.raises(httpx.ReadError): + with pytest.raises(httpx.ConnectError): await ProxyUpdateSpend.update_end_user_spend( n_retry_times=1, prisma_client=mock_prisma_client, @@ -99,6 +99,29 @@ async def test_update_end_user_spend_retries_on_connection_error( assert sleeps == [1.0] +@pytest.mark.asyncio +@pytest.mark.parametrize("ambiguous_error_name", ["ReadTimeout", "ReadError"]) +async def test_update_end_user_spend_does_not_retry_post_send_ambiguous_errors( + mock_prisma_client: Any, ambiguous_error_name: str +) -> None: + """Post-send errors are ambiguous and retrying can double-apply increments + (see DB_RETRY_SAFE_ERROR_TYPES); they must raise on the first attempt.""" + import httpx + + err = getattr(httpx, ambiguous_error_name)("ambiguous") + mock_prisma_client.db.tx = MagicMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + with pytest.raises((httpx.ReadTimeout, httpx.ReadError)): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"u": 1.0}, + ) + mock_prisma_client.db.tx.assert_called_once() + + @pytest.mark.asyncio async def test_update_end_user_spend_non_connection_error_raises_immediately( mock_prisma_client: Any, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a0b3af54750..d9eeb168611 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -188,6 +188,35 @@ async def test_update_spend_logs_job_skips_when_queue_empty( assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 0 +@pytest.mark.asyncio +async def test_update_spend_logs_job_drains_tool_queue_when_spend_queue_empty( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # Regression: a spend-log write failure aborts a run before the tool drain, + # so tool transactions can outlive the spend queue; the job must still run + # for them instead of early-returning on the empty spend queue. + 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() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [MagicMock()] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock() + monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False) + flush_stub = AsyncMock() + monkeypatch.setattr(tool_mod, "flush_tool_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 len(flush_stub.await_args.kwargs["transactions"]) == 1 + assert mock_prisma_client.tool_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 @@ -208,7 +237,7 @@ async def test_update_spend_logs_job_processes_and_clears_queue( guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False ) monkeypatch.setattr( - tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False + tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False ) await update_spend_logs_job( diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d79c755ebdd..1c07f2d6247 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -225,11 +225,6 @@ "count": 1 } }, - "src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx": { "local/filename-pascal-case": { "count": 1 @@ -4337,4 +4332,4 @@ "count": 1 } } -} +} \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 27261768b8d..363525c48af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -5,7 +5,7 @@ const mockUserDailyActivityCall = vi.fn(); vi.mock("@/components/networking", () => ({ userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), - getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }), + getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], start_date: null, end_date: null }), getGeneralSettingsCall: vi.fn().mockResolvedValue([]), })); @@ -19,7 +19,7 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: () =>
, BarChart: () =>
, CustomLegend: () =>
, - DEFAULT_COLOR_CYCLE: ["emerald"], + SEQUENTIAL_COLOR_RAMP: ["indigo"], })); vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index f84167f5a82..5c26ac30477 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -23,18 +23,37 @@ vi.mock("@/components/shared/charts", () => ({ DonutChart: ({ data, label }: { data: unknown; label: string }) => (
), - BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( -
+ BarChart: ({ + data, + categories, + colors, + showLegend, + maxBarSize, + }: { + data: unknown; + categories: string[]; + colors?: readonly string[]; + showLegend?: boolean; + maxBarSize?: number; + }) => ( +
), CustomLegend: ({ categories }: { categories: readonly string[] }) => (
{categories.join(",")}
), - DEFAULT_COLOR_CYCLE: ["emerald", "blue", "violet", "amber"], + SEQUENTIAL_COLOR_RAMP: ["indigo", "blue", "sky", "cyan"], })); import UsageTab from "./UsageTab"; -const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], total_spend: 0, start_date: null, end_date: null }; +const emptyToolSpend: ToolSpendResponse = { by_tool: [], daily: [], start_date: null, end_date: null }; const baseMetrics = (overrides: Partial): SpendMetrics => ({ spend: 0, @@ -216,7 +235,6 @@ describe("UsageTab", () => { { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, ], daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 5.0, start_date: "2026-07-12", end_date: "2026-07-12", }; @@ -225,32 +243,29 @@ describe("UsageTab", () => { const bars = await findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); + // The 64px bar cap is this card's opt-in; the shared BarChart must not cap + // by default (other consumers keep their pre-existing geometry). + expect(bars[0].getAttribute("data-max-bar-size")).toBe("64"); }); - it("notes the 30-day cap when the server clamps the tool spend window", async () => { + it("renders the tool legend once outside the charts, with both charts sharing the tool colors", async () => { const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], + by_tool: [ + { tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }, + { tool_name: "read_file", spend: 1.0, call_count: 2, total_tokens: 50 }, + ], daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-05", - end_date: "2026-07-14", + start_date: "2026-07-12", + end_date: "2026-07-12", }; - const { findByText } = renderWith([day("2026-07-12", {})], { toolSpend }); + const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); - expect(await findByText(/capped at 30 days before the end of the selected range/)).toBeInTheDocument(); - }); + const bars = await findAllByTestId("bar-chart"); + const [totalByTool, dailyByTool] = bars.slice(-2); + expect(dailyByTool.getAttribute("data-show-legend")).toBe("false"); + expect(totalByTool.getAttribute("data-colors")).toBe(dailyByTool.getAttribute("data-colors")); - it("shows no cap note when the served window matches the request", async () => { - const toolSpend = { - by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], - daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], - total_spend: 4.0, - start_date: "2026-07-01", - end_date: "2026-07-14", - }; - const { findAllByTestId, queryByText } = renderWith([day("2026-07-12", {})], { toolSpend }); - - await findAllByTestId("bar-chart"); - expect(queryByText(/capped at 30 days before the end of the selected range/)).not.toBeInTheDocument(); + const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + expect(toolLegends).toHaveLength(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index 15ce84b8445..ec37418e0b5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -3,7 +3,7 @@ import React, { useEffect, useMemo, useState } from "react"; import { Info } from "lucide-react"; -import { AreaChart, BarChart, CustomLegend, DonutChart, DEFAULT_COLOR_CYCLE } from "@/components/shared/charts"; +import { AreaChart, BarChart, CustomLegend, DonutChart, SEQUENTIAL_COLOR_RAMP } from "@/components/shared/charts"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; @@ -34,7 +34,6 @@ interface UsageTabProps { const EMPTY_TOOL_SPEND: ToolSpendResponse = { by_tool: [], daily: [], - total_spend: 0, start_date: null, end_date: null, }; @@ -103,7 +102,6 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { const toolSpend = toolSpendState?.key === rangeKey ? toolSpendState.data : null; const toolSpendLoading = toolSpendEnabled && toolSpend === null; - const toolSpendWindowClamped = !!toolSpend?.start_date && !!startTime && toolSpend.start_date > isoDay(startTime); const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); @@ -168,7 +166,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { })), [toolSpend, topToolNames], ); - const toolColors = useMemo(() => DEFAULT_COLOR_CYCLE.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); + const toolColors = useMemo(() => SEQUENTIAL_COLOR_RAMP.slice(0, Math.max(topToolNames.length, 1)), [topToolNames]); return (
@@ -262,15 +260,10 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { Spend by tool

- Spend on requests that called each tool (MCP and client-side tools). A request that used multiple tools - counts its full spend toward each, so this attributes rather than partitions spend. + Spend on requests that invoked each tool (MCP and client-side tools); declaring a tool without invoking it + does not count. A request that invoked multiple tools counts its full spend toward each, so this attributes + rather than partitions spend.

- {toolSpendWindowClamped && ( -

- Tool spend is capped at 30 days before the end of the selected range; showing spend since{" "} - {toolSpend?.start_date}. -

- )} {topTools.length === 0 ? ( @@ -285,22 +278,27 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { data={topToolsChart} index="tool_name" categories={["spend"]} - colors={["emerald"]} + colors={toolColors} + colorByDatum layout="vertical" yAxisWidth={140} + maxBarSize={64} showLegend={false} valueFormatter={usd} />

Daily spend by tool

+
diff --git a/ui/litellm-dashboard/src/components/ToolDetail.tsx b/ui/litellm-dashboard/src/components/ToolDetail.tsx index 6a8457a559f..06f14638141 100644 --- a/ui/litellm-dashboard/src/components/ToolDetail.tsx +++ b/ui/litellm-dashboard/src/components/ToolDetail.tsx @@ -430,7 +430,7 @@ export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) {

- Recent logs + Recent invocations

{ expect(container.querySelector("style")).toBeNull(); }); + it("colors each bar by its datum when colorByDatum is set, instead of one fill for the series", () => { + const singleCategory = [ + { tool: "alpha", spend: 3 }, + { tool: "beta", spend: 2 }, + { tool: "gamma", spend: 1 }, + ]; + + const { container, rerender } = render( + , + ); + const sharedFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(new Set(sharedFills).size).toBe(1); + + rerender( + , + ); + const perDatumFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) => + rect.getAttribute("fill"), + ); + expect(perDatumFills).toEqual([ + "var(--color-blue-500, #3b82f6)", + "var(--color-cyan-500, #06b6d4)", + "var(--color-violet-500, #8b5cf6)", + ]); + }); + it("stacks bars into a single column per index when stack is set", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx index 7069ececb70..ab2cc66eaf4 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx @@ -1,7 +1,7 @@ "use client"; import * as React from "react"; -import { Bar, BarChart as RechartsBarChart, CartesianGrid, XAxis, YAxis } from "recharts"; +import { Bar, BarChart as RechartsBarChart, CartesianGrid, Cell, XAxis, YAxis } from "recharts"; import { ChartContainer, ChartLegend, ChartLegendContent, ChartTooltip, type ChartConfig } from "@/components/ui/chart"; import { cn } from "@/lib/cva.config"; import { ValueTooltip, type ChartTooltipComponent } from "./chart_tooltip"; @@ -12,6 +12,8 @@ export type BarChartProps> = { index: string; categories: readonly string[]; colors?: readonly ChartColor[]; + colorByDatum?: boolean; + maxBarSize?: number; valueFormatter?: (value: number) => string; stack?: boolean; layout?: "horizontal" | "vertical"; @@ -32,6 +34,8 @@ export function BarChart>({ index, categories, colors, + colorByDatum = false, + maxBarSize, valueFormatter, stack = false, layout = "horizontal", @@ -57,7 +61,7 @@ export function BarChart>({ ); } - const fills = categoryFills(categories.length, colors); + const fills = categoryFills(colorByDatum ? data.length : categories.length, colors); const config: ChartConfig = Object.fromEntries(categories.map((category) => [category, { label: category }])); const vertical = layout === "vertical"; const TooltipContent = customTooltip ?? ValueTooltip; @@ -115,6 +119,7 @@ export function BarChart>({ fill={fills[i]} stackId={stack ? "stack" : undefined} isAnimationActive={false} + maxBarSize={maxBarSize} onClick={ onValueChange ? (item: { payload?: TDatum }) => { @@ -122,7 +127,9 @@ export function BarChart>({ } : undefined } - /> + > + {colorByDatum && data.map((_, dataIndex) => )} + ))} diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx index 889927aca43..28afe5faf9c 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx @@ -21,6 +21,14 @@ describe("CustomLegend", () => { expect(dots[1]?.getAttribute("style")).toContain("--color-green-500"); }); + it("wraps onto multiple lines instead of overflowing when there are many categories", () => { + const { container } = render( + `metrics.tool_${i}`)} colors={["blue", "green"]} />, + ); + + expect(container.firstElementChild?.className).toContain("flex-wrap"); + }); + it("cycles colors when there are more categories than colors", () => { const { container } = render( , diff --git a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx index da252d8bf63..1551f3d0e39 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx @@ -11,7 +11,7 @@ export const CustomLegend = ({ categories: readonly string[]; colors: readonly ChartColor[]; }) => ( -
+
{categories.map((category, idx) => (
`var(--color-${color}-500, ${CHART_COLOR_HEX[color]})`; +export const SEQUENTIAL_COLOR_RAMP: readonly ChartColor[] = [ + "#1e3a8a", + "#1d4ed8", + "#2563eb", + "#3b82f6", + "#60a5fa", + "#93c5fd", + "#bfdbfe", + "#dbeafe", +]; + +const NAMED_COLOR_HEX: Readonly> = CHART_COLOR_HEX; + +export const chartColorValue = (color: ChartColor): string => + color in NAMED_COLOR_HEX ? `var(--color-${color}-500, ${NAMED_COLOR_HEX[color]})` : color; export const categoryFills = (count: number, colors?: readonly ChartColor[]): readonly string[] => { const cycle = colors && colors.length > 0 ? colors : DEFAULT_COLOR_CYCLE; diff --git a/ui/litellm-dashboard/src/components/shared/charts/index.ts b/ui/litellm-dashboard/src/components/shared/charts/index.ts index 8383c767064..69edd3fb13f 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/index.ts +++ b/ui/litellm-dashboard/src/components/shared/charts/index.ts @@ -8,6 +8,13 @@ export { type ChartTooltipComponent, type ChartTooltipProps, } from "./chart_tooltip"; -export { CHART_COLOR_HEX, DEFAULT_COLOR_CYCLE, categoryFills, chartColorValue, type ChartColor } from "./colors"; +export { + CHART_COLOR_HEX, + DEFAULT_COLOR_CYCLE, + SEQUENTIAL_COLOR_RAMP, + categoryFills, + chartColorValue, + type ChartColor, +} from "./colors"; export { DonutChart, type DonutChartProps } from "./donut_chart"; export { LineChart, type LineChartCurveType, type LineChartProps } from "./line_chart"; diff --git a/ui/litellm-dashboard/src/components/ui/chart.tsx b/ui/litellm-dashboard/src/components/ui/chart.tsx index 14e10b9f06f..f0310819097 100644 --- a/ui/litellm-dashboard/src/components/ui/chart.tsx +++ b/ui/litellm-dashboard/src/components/ui/chart.tsx @@ -264,7 +264,11 @@ const ChartLegendContent = React.forwardRef< return (
{payload .filter((item) => item.type !== "none") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fc08864f19f..541d7a17ae1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -18005,16 +18005,16 @@ export interface paths { * Get Tool Spend * @description Spend attributed to each tool over a date range, for the Cost Optimization dashboard. * - * Joins ``LiteLLM_SpendLogToolIndex`` (which tool names ran on which request) to - * ``LiteLLM_SpendLogs`` (what the request cost). A request that used multiple tools - * counts its full spend toward each of those tools, so per-tool numbers are - * attributions. ``total_spend`` is the deduplicated spend of every request that - * called at least one tool in the window, so it never double counts. + * Reads the ``LiteLLM_DailyToolSpend`` rollup, written at request time from invoked + * tools only (MCP tool calls and response tool_calls; declaring a tool without + * invoking it does not count). A request that invoked multiple tools counts its + * full spend toward each of them, so per-tool numbers are attributions and do not + * sum to a deduplicated total. * - * ``start_date`` is clamped to at most 30 days before ``end_date`` (serving up to - * 31 calendar dates inclusive, the same width as the endpoint's default window): - * a wider requested range is clamped, and the response's ``start_date`` reflects - * the effective window actually served. + * ``by_tool`` is the top ``TOOL_SPEND_TOP_TOOLS`` tools by spend, aggregated in + * SQL, and ``daily`` covers only those tools, so the response is bounded by + * days x TOOL_SPEND_TOP_TOOLS regardless of the requested range or how many + * distinct tool names exist. */ get: operations["get_tool_spend_v1_tool_spend_get"]; put?: never; @@ -18074,7 +18074,8 @@ export interface paths { }; /** * Get Tool Usage Logs - * @description Return paginated spend logs for requests that used this tool (from SpendLogToolIndex). + * @description Return paginated spend logs for requests that invoked this tool (from SpendLogToolIndex). + * Declaring a tool in a request body without the model invoking it does not create an entry. */ get: operations["get_tool_usage_logs_v1_tool__tool_name__logs_get"]; put?: never; @@ -32211,12 +32212,6 @@ export interface components { end_date?: string | null; /** Start Date */ start_date?: string | null; - /** - * Total Spend - * @description Deduplicated spend of every request that called at least one tool in the window; less than the sum of per-tool attributed spend whenever multi-tool requests exist - * @default 0 - */ - total_spend: number; }; /** * ToolUsageLogEntry