scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #34675 from BerriAI/litellm_tool_spend_rollup

Browse source 
fix(proxy): roll up tool spend daily instead of scanning SpendLogs
This commit is contained in:
tin-berri 2026-07-27 15:31:19 -07:00 • committed by GitHub
commit 9bb75d67af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 1348 additions and 494 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
db_scripts/backfill_daily_tool_spend.sql Normal file
View file

@ -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;

12
litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql Normal file
View file

@ -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")
);

13
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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())

2
litellm/constants.py
View file

@ -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))

32
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -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,
):

10
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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": [
{

7
litellm/proxy/_types.py
View file

@ -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]

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

@ -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,

264
litellm/proxy/db/spend_log_tool_index.py
View file

@ -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))

170
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

@ -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

13
litellm/proxy/schema.prisma
View file

@ -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())

45
litellm/proxy/utils.py
View file

@ -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:

2
litellm/repositories/__init__.py
View file

@ -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",

4
litellm/repositories/table_repositories.py
View file

@ -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"

7
litellm/types/tool_management.py
View file

@ -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

13
schema.prisma
View file

@ -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())

4
tests/proxy_unit_tests/test_update_spend.py
View file

@ -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

31
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -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"]

234
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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():
"""

348
tests/test_litellm/proxy/db/test_spend_log_tool_index.py Normal file
View file

@ -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()

220
tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py
View file

@ -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()

6
tests/test_litellm/proxy/test_spend_log_cleanup.py
View file

@ -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():

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

@ -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()

35
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -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,

31
tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
View file

@ -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(

7
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -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
}
}
}
}

4
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx
View file

@ -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: () => <div />,
BarChart: () => <div />,
CustomLegend: () => <div />,
DEFAULT_COLOR_CYCLE: ["emerald"],
SEQUENTIAL_COLOR_RAMP: ["indigo"],
}));
vi.mock("@/app/(dashboard)/router-settings/_components/general_settings", () => ({

65
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx
View file

@ -23,18 +23,37 @@ vi.mock("@/components/shared/charts", () => ({
DonutChart: ({ data, label }: { data: unknown; label: string }) => (
<div data-testid="donut-chart" data-label={label} data-slices={JSON.stringify(data)} />
),
BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => (
<div data-testid="bar-chart" data-categories={categories.join(",")} data-series={JSON.stringify(data)} />
BarChart: ({
data,
categories,
colors,
showLegend,
maxBarSize,
}: {
data: unknown;
categories: string[];
colors?: readonly string[];
showLegend?: boolean;
maxBarSize?: number;
}) => (
<div
data-testid="bar-chart"
data-categories={categories.join(",")}
data-colors={(colors ?? []).join(",")}
data-show-legend={String(showLegend ?? true)}
data-max-bar-size={maxBarSize === undefined ? "" : String(maxBarSize)}
data-series={JSON.stringify(data)}
/>
),
CustomLegend: ({ categories }: { categories: readonly string[] }) => (
<div data-testid="chart-legend">{categories.join(",")}</div>
),
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>): 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);
});
});

24
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx
View file

@ -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<UsageTabProps> = ({ 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<UsageTabProps> = ({ 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 (
<div className="w-full space-y-6">
@ -262,15 +260,10 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
<CardHeader>
<CardTitle>Spend by tool</CardTitle>
<p className="text-sm text-muted-foreground">
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.
</p>
{toolSpendWindowClamped && (
<p className="text-xs text-muted-foreground">
Tool spend is capped at 30 days before the end of the selected range; showing spend since{" "}
{toolSpend?.start_date}.
</p>
)}
</CardHeader>
<CardContent>
{topTools.length === 0 ? (
@ -285,22 +278,27 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
data={topToolsChart}
index="tool_name"
categories={["spend"]}
colors={["emerald"]}
colors={toolColors}
colorByDatum
layout="vertical"
yAxisWidth={140}
maxBarSize={64}
showLegend={false}
valueFormatter={usd}
/>
</div>
<div>
<p className="mb-2 text-sm font-medium text-muted-foreground">Daily spend by tool</p>
<CustomLegend categories={topToolNames} colors={toolColors} />
<BarChart
data={dailyToolSeries}
index="date"
categories={topToolNames}
colors={toolColors}
stack
maxBarSize={64}
valueFormatter={usd}
showLegend={false}
/>
</div>
</div>

2
ui/litellm-dashboard/src/components/ToolDetail.tsx
View file

@ -430,7 +430,7 @@ export function ToolDetail({ toolName, onBack, accessToken }: ToolDetailProps) {
<section className="rounded-lg border border-border bg-card p-5 shadow-xs">
<h2 className="mb-3 flex items-center gap-2 text-sm font-semibold">
<History className="size-4" />
Recent logs
Recent invocations
</h2>
<LogViewer
guardrailName={tool.tool_name}

1
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -7527,7 +7527,6 @@ export interface ToolSpendDailyEntry {
export interface ToolSpendResponse {
by_tool: ToolSpendEntry[];
daily: ToolSpendDailyEntry[];
total_spend: number;
start_date: string | null;
end_date: string | null;
}

34
ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx
View file

@ -113,6 +113,40 @@ describe("BarChart", () => {
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(
<BarChart data={singleCategory} index="tool" categories={["spend"]} colors={["blue", "cyan", "violet"]} />,
);
const sharedFills = Array.from(container.querySelectorAll("path.recharts-rectangle")).map((rect) =>
rect.getAttribute("fill"),
);
expect(new Set(sharedFills).size).toBe(1);
rerender(
<BarChart
data={singleCategory}
index="tool"
categories={["spend"]}
colors={["blue", "cyan", "violet"]}
colorByDatum
/>,
);
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(
<BarChart data={data} index="date" categories={["passed", "blocked"]} colors={["green", "red"]} stack={true} />,

13
ui/litellm-dashboard/src/components/shared/charts/bar_chart.tsx
View file

@ -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<TDatum extends Record<string, unknown>> = {
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<TDatum extends Record<string, unknown>>({
index,
categories,
colors,
colorByDatum = false,
maxBarSize,
valueFormatter,
stack = false,
layout = "horizontal",
@ -57,7 +61,7 @@ export function BarChart<TDatum extends Record<string, unknown>>({
);
}
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<TDatum extends Record<string, unknown>>({
fill={fills[i]}
stackId={stack ? "stack" : undefined}
isAnimationActive={false}
maxBarSize={maxBarSize}
onClick={
onValueChange
? (item: { payload?: TDatum }) => {
@ -122,7 +127,9 @@ export function BarChart<TDatum extends Record<string, unknown>>({
}
: undefined
}
/>
>
{colorByDatum && data.map((_, dataIndex) => <Cell key={dataIndex} fill={fills[dataIndex]} />)}
</Bar>
))}
</RechartsBarChart>
</ChartContainer>

8
ui/litellm-dashboard/src/components/shared/charts/chart_legend.test.tsx
View file

@ -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(
<CustomLegend categories={Array.from({ length: 8 }, (_, i) => `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(
<CustomLegend categories={["metrics.a", "metrics.b", "metrics.c"]} colors={["blue", "green"]} />,

2
ui/litellm-dashboard/src/components/shared/charts/chart_legend.tsx
View file

@ -11,7 +11,7 @@ export const CustomLegend = ({
categories: readonly string[];
colors: readonly ChartColor[];
}) => (
<div className="flex items-center justify-end space-x-4">
<div className="flex flex-wrap items-center justify-end gap-x-4 gap-y-1">
{categories.map((category, idx) => (
<div key={category} className="flex items-center space-x-2">
<span

18
ui/litellm-dashboard/src/components/shared/charts/colors.ts
View file

@ -23,7 +23,7 @@ export const CHART_COLOR_HEX = {
rose: "#f43f5e",
} as const;
export type ChartColor = keyof typeof CHART_COLOR_HEX;
export type ChartColor = keyof typeof CHART_COLOR_HEX | `#${string}`;
export const DEFAULT_COLOR_CYCLE: readonly ChartColor[] = [
"blue",
@ -50,7 +50,21 @@ export const DEFAULT_COLOR_CYCLE: readonly ChartColor[] = [
"rose",
];
export const chartColorValue = (color: ChartColor): string => `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<Record<string, string>> = 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;

9
ui/litellm-dashboard/src/components/shared/charts/index.ts
View file

@ -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";

6
ui/litellm-dashboard/src/components/ui/chart.tsx
View file

@ -264,7 +264,11 @@ const ChartLegendContent = React.forwardRef<
return (
<div
ref={ref}
className={cn("flex items-center justify-center gap-4", verticalAlign === "top" ? "pb-3" : "pt-3", className)}
className={cn(
"flex flex-wrap items-center justify-center gap-x-4 gap-y-1",
verticalAlign === "top" ? "pb-3" : "pt-3",
className,
)}
>
{payload
.filter((item) => item.type !== "none")

27
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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