scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/budgets-table-truncation-4b98ec

Browse source
This commit is contained in:
Yuneng Jiang 2026-07-30 14:04:18 -07:00
commit 683d716ca2
No known key found for this signature in database
17 changed files with 381 additions and 188 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/a2a_protocol/main.py
View file

@ -568,6 +568,7 @@ def _build_streaming_logging_obj(
logging_obj.custom_llm_provider = "a2a_agent"
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
logging_obj.model_call_details["call_type"] = logging_obj.call_type
if agent_id:
logging_obj.model_call_details["agent_id"] = agent_id

1
litellm/integrations/opentelemetry.py
View file

@ -118,6 +118,7 @@ TOKEN_TYPE_ATTRIBUTE: str = "gen_ai.token.type"
VALID_METRIC_ATTRIBUTE_NAMES: FrozenSet[str] = frozenset(
(
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",

72
litellm/integrations/otel/model/semconv.py
View file

@ -6,17 +6,30 @@ without a semconv equivalent lives under the ``litellm.*`` vendor namespace.
from enum import Enum
from typing import Final
from litellm._logging import verbose_logger
class GenAIOperation(str, Enum):
"""Values for ``gen_ai.operation.name``."""
"""Values for ``gen_ai.operation.name``.
The first block is the convention's own vocabulary. The ``LITELLM_`` members
are vendor values for operations the convention names nothing for; its note
on this attribute directs instrumentation to use a system-specific name in
exactly that case, the same allowance :func:`resolve_provider` relies on for
unmapped providers. They stay under the ``litellm.`` prefix so a value the
convention adds later can never collide with one of ours.
"""
CHAT = "chat"
TEXT_COMPLETION = "text_completion"
EMBEDDINGS = "embeddings"
GENERATE_CONTENT = "generate_content"
RETRIEVAL = "retrieval" # vector-store search / RAG query spans
CREATE_AGENT = "create_agent" # reserved for future agent spans
INVOKE_AGENT = "invoke_agent" # reserved for future agent spans
INVOKE_AGENT = "invoke_agent" # agent (A2A) message spans
EXECUTE_TOOL = "execute_tool" # MCP tool-call spans
LITELLM_VECTOR_STORE_MANAGEMENT = "litellm.vector_store_management"
LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management"
class GenAIProvider(str, Enum):
@ -49,11 +62,17 @@ class MCPMethod(str, Enum):
class GenAI:
"""Canonical OTel GenAI span-attribute keys."""
"""Canonical OTel GenAI attribute keys.
``SYSTEM`` is the one exception: the convention deprecated it in favor of
``PROVIDER_NAME``, and it survives here only so already-shipped series keep
resolving for consumers that query it. Nothing new should use it.
"""
# request
OPERATION_NAME: Final = "gen_ai.operation.name"
PROVIDER_NAME: Final = "gen_ai.provider.name"
SYSTEM: Final = "gen_ai.system"
REQUEST_MODEL: Final = "gen_ai.request.model"
REQUEST_TEMPERATURE: Final = "gen_ai.request.temperature"
REQUEST_TOP_P: Final = "gen_ai.request.top_p"
@ -316,6 +335,35 @@ _OPERATION_BY_CALL_TYPE: dict[str, GenAIOperation] = {
"responses": GenAIOperation.CHAT,
"aresponses": GenAIOperation.CHAT,
"call_mcp_tool": GenAIOperation.EXECUTE_TOOL,
"vector_store_search": GenAIOperation.RETRIEVAL,
"avector_store_search": GenAIOperation.RETRIEVAL,
"query": GenAIOperation.RETRIEVAL,
"aquery": GenAIOperation.RETRIEVAL,
"send_message": GenAIOperation.INVOKE_AGENT,
"asend_message": GenAIOperation.INVOKE_AGENT,
"asend_message_streaming": GenAIOperation.INVOKE_AGENT,
"vector_store_create": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_create": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_list": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_list": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_update": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_update": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_delete": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_delete": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_file_create": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_create": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_list": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_list": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_content": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_content": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_update": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_update": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_delete": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_delete": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
}
@ -332,7 +380,21 @@ def resolve_provider(custom_llm_provider: str | None) -> str:
def resolve_operation(call_type: str | None) -> GenAIOperation:
"""Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value."""
"""Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value.
An unmapped call type still falls back to ``chat`` so every series keeps an
operation label, but it logs at debug rather than falling through silently:
a new call type mislabelled as ``chat`` mixes its latency and cost into
everyone's chat charts, which is invisible until someone reads the numbers.
"""
if not call_type:
return GenAIOperation.CHAT
return _OPERATION_BY_CALL_TYPE.get(call_type.lower(), GenAIOperation.CHAT)
mapped = _OPERATION_BY_CALL_TYPE.get(call_type.lower())
if mapped is not None:
return mapped
verbose_logger.debug(
"otel: call_type %r has no gen_ai.operation.name mapping; labelling it %r. Add it to _OPERATION_BY_CALL_TYPE.",
call_type,
GenAIOperation.CHAT.value,
)
return GenAIOperation.CHAT

33
litellm/integrations/otel/plumbing/metrics.py
View file

@ -23,11 +23,34 @@ from litellm.integrations.opentelemetry import (
_resolve_metric_attribute_filter,
)
from litellm.integrations.otel.model.metadata import time_to_first_chunk_seconds
from litellm.integrations.otel.model.semconv import Error, Metric, resolve_operation
from litellm.integrations.otel.model.semconv import (
Error,
GenAI,
Metric,
resolve_operation,
resolve_provider,
)
from litellm.integrations.otel.model.utils import to_seconds
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
def _provider_attributes(custom_llm_provider: object) -> Mapping[str, str]:
"""The provider labels for one call's metrics.
``gen_ai.provider.name`` carries the semconv-mapped value; the deprecated
``gen_ai.system`` spelling is dual-emitted with the raw litellm provider
string it has always carried, so a dashboard already querying it keeps
matching. A call with no provider gets neither label: a placeholder value
would mint a permanent series that no operator can act on.
"""
if not isinstance(custom_llm_provider, str) or not custom_llm_provider:
return {}
return {
GenAI.PROVIDER_NAME: resolve_provider(custom_llm_provider),
GenAI.SYSTEM: custom_llm_provider,
}
@dataclass(frozen=True)
class GenAIMetrics:
operation_duration: Histogram
@ -98,6 +121,7 @@ ERROR_TYPE_FALLBACK: Final = "_OTHER"
METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset(
(
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",
@ -223,11 +247,10 @@ class GenAIMetricRecorder:
def _common_attributes(self, kwargs: Mapping[str, Any]) -> dict:
params = kwargs.get("litellm_params") or {}
provider = params.get("custom_llm_provider", "Unknown")
common_attrs: dict = {
"gen_ai.operation.name": resolve_operation(kwargs.get("call_type")).value,
"gen_ai.system": provider,
"gen_ai.request.model": kwargs.get("model"),
GenAI.OPERATION_NAME: resolve_operation(kwargs.get("call_type")).value,
**_provider_attributes(params.get("custom_llm_provider")),
GenAI.REQUEST_MODEL: kwargs.get("model"),
"gen_ai.framework": "litellm",
}

10
litellm/proxy/_types.py
View file

@ -2461,16 +2461,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"is active as a reminder that hard enforcement is relaxed."
),
)
skip_user_budget_on_team_key: bool | None = Field(
None,
description=(
"If True, restores the legacy behavior where a user's personal "
"max_budget is NOT enforced when their key belongs to a team; only "
"the team (and team-member) budgets apply. Defaults to False, meaning "
"the user's personal max_budget is always enforced regardless of "
"whether the key belongs to a team (see GitHub issue #12905)."
),
)
user_url_validation: Optional[bool] = Field(
None,
description=(

41
litellm/proxy/auth/auth_checks.py
View file

@ -633,31 +633,28 @@ async def common_checks(
)
async def _user_max_budget_check() -> None:
if user_object is None or user_object.max_budget is None:
return
skip_for_team = (
general_settings.get("skip_user_budget_on_team_key") is True
and team_object is not None
and team_object.team_id is not None
)
if skip_for_team:
return
from litellm.proxy.proxy_server import get_current_spend
# 4.1 personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
from litellm.proxy.proxy_server import get_current_spend
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
)
# Each scope reads a distinct counter key with no cross-scope ordering
# dependency, so the per-scope Redis-first reads run concurrently instead

1
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2450,7 +2450,6 @@ async def _reserve_budget_after_common_checks(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)

1
litellm/proxy/proxy_server.py
View file

@ -15258,7 +15258,6 @@ async def get_config_list(
"forward_client_headers_to_llm_api": {"type": "Boolean"},
"mcp_required_fields": {"type": "List"},
"cancel_on_disconnect": {"type": "Boolean"},
"skip_user_budget_on_team_key": {"type": "Boolean"},
"disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"},
}

6
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -155,7 +155,6 @@ async def reserve_budget_for_request(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
fail_closed_budget_enforcement: bool = False,
) -> Optional[dict]:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
@ -175,7 +174,6 @@ async def reserve_budget_for_request(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=skip_user_budget_on_team_key,
)
if not counters:
return None
@ -333,7 +331,6 @@ async def _get_budget_counters(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
) -> List[_BudgetCounter]:
counters: List[_BudgetCounter] = []
@ -382,9 +379,8 @@ async def _get_budget_counters(
)
)
is_team_key = team_object is not None and team_object.team_id is not None
if (
not (is_team_key and skip_user_budget_on_team_key)
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.user_id is not None
and user_object.max_budget is not None

11
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -219,8 +219,8 @@ async def test_aaauser_personal_budgets(key_ownership):
"""
Set a personal budget on a user
User budget is enforced regardless of key ownership (personal or team).
Both cases should raise BudgetExceededError when the user is over budget.
- have it only apply when key belongs to user -> raises BudgetExceededError
- if key belongs to team, have key respect team budget -> allows call to go through
"""
import asyncio
import time
@ -278,9 +278,12 @@ async def test_aaauser_personal_budgets(key_ownership):
== valid_token
)
with pytest.raises(ProxyException) as exc_info:
if key_ownership == "user_key":
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
else:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
@pytest.mark.asyncio

29
tests/test_litellm/a2a_protocol/test_main.py
View file

@ -104,3 +104,32 @@ async def test_streaming_trace_id_prefers_logging_trace_id():
pass
assert captured["extra_headers"]["X-LiteLLM-Trace-Id"] == "trace-from-logging"
def test_streaming_logging_obj_carries_call_type_into_model_call_details():
"""The streaming logging object is built by hand rather than through
``update_environment_variables``, which is the only place ``call_type`` normally
reaches ``model_call_details``. Callbacks read the call type from there, so
without this the streamed turn arrives at every logger with no call type at all
and OTel's GenAI metrics label it ``chat`` instead of ``invoke_agent``."""
from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest
from litellm.a2a_protocol.main import _build_streaming_logging_obj
request = SendStreamingMessageRequest(
id="rpc-call-type",
params=MessageSendParams(
message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
),
)
logging_obj = _build_streaming_logging_obj(
request=request,
agent_name="some-agent",
agent_id=None,
litellm_params=None,
metadata=None,
proxy_server_request=None,
)
assert logging_obj.model_call_details["call_type"] == "asend_message_streaming"

117
tests/test_litellm/integrations/otel/test_otel_v2_metrics.py
View file

@ -68,6 +68,9 @@ ALL_METRICS = frozenset(
TOKEN_TYPE = "gen_ai.token.type"
MODEL_KEY = "gen_ai.request.model"
OPERATION_KEY = "gen_ai.operation.name"
PROVIDER_NAME_KEY = "gen_ai.provider.name"
SYSTEM_KEY = "gen_ai.system"
# Keys inside the ceiling that an operator's filter must still be able to remove.
# Every one is bounded, so it survives the ceiling and only the operator's own
@ -83,18 +86,25 @@ COMPLETION_TOKENS = 89
RESPONSE_COST = 0.0023
def _build_call(stream: bool = True):
def _build_call(
stream: bool = True,
provider: str | None = "openai",
call_type: str = "completion",
):
"""A captured success-call (kwargs, response_obj, start, end) that exercises
every one of the six metrics: usage for token.usage, response_cost for cost,
streaming + timing for the response-time histograms."""
streaming + timing for the response-time histograms.
``provider=None`` omits ``custom_llm_provider`` entirely, reproducing a call
litellm could not attribute to a provider."""
start = datetime(2026, 6, 12, 12, 0, 0)
api_call_start = start + timedelta(seconds=0.1)
completion_start = start + timedelta(seconds=0.5)
end = start + timedelta(seconds=1.0)
kwargs = {
"model": "gpt-4o-mini",
"call_type": "completion",
"litellm_params": {"custom_llm_provider": "openai"},
"call_type": call_type,
"litellm_params": ({"custom_llm_provider": provider} if provider is not None else {}),
"optional_params": {"stream": stream},
"response_cost": RESPONSE_COST,
"api_call_start_time": api_call_start,
@ -142,7 +152,7 @@ def _metrics_by_name(reader):
return out
def _drive_success(reader, callback_settings_attributes=None):
def _drive_success(reader, callback_settings_attributes=None, **call_overrides):
"""Construct a metrics-on logger, optionally populate callback_settings AFTER
construction (mirroring the proxy ordering), run the real success hook."""
logger = _logger(reader, enable_metrics=True)
@ -152,7 +162,7 @@ def _drive_success(reader, callback_settings_attributes=None):
"otel": {"attributes": callback_settings_attributes}
}
try:
kwargs, response_obj, start, end = _build_call()
kwargs, response_obj, start, end = _build_call(**call_overrides)
asyncio.run(logger.async_log_success_event(kwargs, response_obj, start, end))
finally:
litellm.callback_settings = previous
@ -376,6 +386,100 @@ def test_success_attributes_are_capped_at_the_ceiling():
for dp in metrics[name]:
leaked = set(dp.attributes) - set(BOUNDED_KEYS) - {TOKEN_TYPE}
assert not leaked, f"{name} leaked {leaked}"
def test_provider_is_labelled_with_semconv_provider_name():
"""Every recorded point carries gen_ai.provider.name holding the semconv
provider value (bedrock -> aws.bedrock), the key the GenAI convention and the
dashboards built on it query. The deprecated gen_ai.system spelling alone is
unreadable to them."""
metrics = _drive_success(InMemoryMetricReader(), provider="bedrock")
for name in ALL_METRICS:
points = metrics[name]
assert points, f"{name} was not recorded"
for dp in points:
assert dp.attributes[PROVIDER_NAME_KEY] == "aws.bedrock"
def test_deprecated_gen_ai_system_is_dual_emitted_verbatim():
"""gen_ai.system keeps its raw litellm provider value alongside the new key
for one release, so a dashboard already filtering on it keeps matching. Its
value must not be swapped for the mapped one, which would break exactly the
queries the dual emission exists to protect."""
metrics = _drive_success(InMemoryMetricReader(), provider="bedrock")
for dp in metrics[OPERATION_DURATION]:
assert dp.attributes[SYSTEM_KEY] == "bedrock"
assert dp.attributes[PROVIDER_NAME_KEY] == "aws.bedrock"
def test_no_provider_attribute_when_provider_is_absent():
"""A call litellm could not attribute to a provider carries no provider label
at all. A placeholder value ("Unknown") would mint a permanent series that
aggregates every unattributable request and that no operator can act on."""
metrics = _drive_success(InMemoryMetricReader(), provider=None)
for name in ALL_METRICS:
points = metrics[name]
assert points, f"{name} was not recorded"
for dp in points:
keys = set(dp.attributes.keys())
assert PROVIDER_NAME_KEY not in keys
assert SYSTEM_KEY not in keys
assert "Unknown" not in set(dp.attributes.values())
def test_vector_store_search_is_not_labelled_as_chat():
"""A vector-store search records under gen_ai.operation.name=retrieval, so its
latency and cost stay out of the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type="avector_store_search")
for name in (OPERATION_DURATION, TOKEN_COST):
for dp in metrics[name]:
assert dp.attributes[OPERATION_KEY] == "retrieval"
@pytest.mark.parametrize(
"call_type,expected",
[
("avector_store_create", "litellm.vector_store_management"),
("avector_store_delete", "litellm.vector_store_management"),
("avector_store_file_create", "litellm.vector_store_file_management"),
("avector_store_file_list", "litellm.vector_store_file_management"),
],
)
def test_vector_store_management_is_not_labelled_as_chat(call_type, expected):
"""Store and file management reach the recorder through the same success hook as a
completion, so leaving them unmapped kept billing- and latency-relevant admin calls
inside the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type=call_type)
for dp in metrics[OPERATION_DURATION]:
assert dp.attributes[OPERATION_KEY] == expected
@pytest.mark.parametrize("call_type", ["asend_message", "asend_message_streaming"])
def test_agent_message_is_not_labelled_as_chat(call_type):
"""An A2A agent send records under gen_ai.operation.name=invoke_agent, streamed or
not. The streaming iterator dispatches the same success handlers under its own
``asend_message_streaming`` call type, so an unmapped streaming spelling puts every
streamed agent turn's latency and cost back into the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type=call_type)
for name in (OPERATION_DURATION, TOKEN_COST):
for dp in metrics[name]:
assert dp.attributes[OPERATION_KEY] == "invoke_agent"
def test_provider_name_is_filterable():
"""gen_ai.provider.name is a member of the metric-attribute allowlist, so an
operator can include or exclude it; an unlisted name raises instead."""
metrics = _drive_success(
InMemoryMetricReader(),
callback_settings_attributes={"include_list": [PROVIDER_NAME_KEY]},
)
for dp in metrics[OPERATION_DURATION]:
assert set(dp.attributes.keys()) == {PROVIDER_NAME_KEY}
def test_metrics_reach_operator_configured_global_provider(monkeypatch):
@ -481,6 +585,7 @@ UNBOUNDED_KEYS = (
BOUNDED_KEYS = (
"hidden_params",
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",

102
tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -1,8 +1,13 @@
"""Tests for the OTel v2 sources of truth: span registry, semconv keys, config,
and the typed StandardLoggingPayload adapter. These need no OTel SDK."""
import logging
import re
from pathlib import Path
import pytest
import litellm
from litellm.integrations.otel import (
BAGGAGE_PROMOTED_KEYS,
DB,
@ -208,6 +213,103 @@ def test_operation_resolution():
assert resolve_operation("call_mcp_tool") is GenAIOperation.EXECUTE_TOOL
@pytest.mark.parametrize("call_type", ["vector_store_search", "avector_store_search"])
def test_vector_store_search_is_a_retrieval_operation(call_type):
"""A vector-store search is a retrieval, so its duration and cost must not
land in the chat series that dashboards read latency off."""
assert resolve_operation(call_type) is GenAIOperation.RETRIEVAL
assert resolve_operation(call_type).value == "retrieval"
@pytest.mark.parametrize("call_type", ["query", "aquery"])
def test_rag_query_is_a_retrieval_operation(call_type):
"""``/rag/query`` reaches the same recorder as a vector-store search and is the
same operation, so it must not be the one retrieval surface left reading as chat."""
assert resolve_operation(call_type) is GenAIOperation.RETRIEVAL
@pytest.mark.parametrize(
"call_type",
[
f"{prefix}vector_store_{verb}"
for verb in ("create", "retrieve", "list", "update", "delete")
for prefix in ("", "a")
],
)
def test_vector_store_management_is_not_chat(call_type):
"""The store lifecycle calls are not GenAI client operations and the convention
names nothing for them, so they take a vendor value rather than defaulting into
the chat series."""
assert resolve_operation(call_type) is GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT
assert resolve_operation(call_type).value == "litellm.vector_store_management"
@pytest.mark.parametrize(
"call_type",
[
f"{prefix}vector_store_file_{verb}"
for verb in ("create", "list", "retrieve", "content", "update", "delete")
for prefix in ("", "a")
],
)
def test_vector_store_file_management_is_not_chat(call_type):
"""The file operations are a distinct REST resource from the store lifecycle, so
they get their own vendor value instead of sharing one bucket."""
assert resolve_operation(call_type) is GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT
assert resolve_operation(call_type).value == "litellm.vector_store_file_management"
def test_vendor_operation_values_are_namespaced():
"""A vendor value must stay under the ``litellm.`` prefix: an unprefixed invented
name could collide with a value the convention adds later, silently changing what
a conformant consumer thinks it is reading."""
vendor = [op for op in GenAIOperation if op.name.startswith("LITELLM_")]
assert vendor, "no vendor operation values defined"
assert all(op.value.startswith("litellm.") for op in vendor)
@pytest.mark.parametrize("call_type", ["send_message", "asend_message", "asend_message_streaming"])
def test_agent_message_is_an_invoke_agent_operation(call_type):
"""An agent (A2A) message send is an agent invocation, not a chat completion.
The streaming spelling counts: ``_build_streaming_logging_obj`` in
``litellm/a2a_protocol/main.py`` stamps ``asend_message_streaming`` on the
logging object the streaming iterator dispatches success handlers with, so a
missing entry sends every streamed agent turn into the chat series. There is
no sync spelling because A2A streaming is async-only.
"""
assert resolve_operation(call_type) is GenAIOperation.INVOKE_AGENT
assert resolve_operation(call_type).value == "invoke_agent"
def test_every_call_type_the_a2a_package_stamps_is_an_agent_operation():
"""Pins the map to the call types the A2A code actually stamps on its logging
objects. A new spelling added there without a map entry fails here instead of
quietly landing in the chat series, which is how the streaming one was missed."""
a2a_package = Path(litellm.__file__).parent / "a2a_protocol"
stamped = {
call_type
for source in a2a_package.rglob("*.py")
for call_type in re.findall(r'call_type="([^"]+)"', source.read_text())
}
assert stamped, "no call_type literals found in litellm/a2a_protocol"
unmapped = {
call_type: resolve_operation(call_type).value
for call_type in stamped
if resolve_operation(call_type) is not GenAIOperation.INVOKE_AGENT
}
assert not unmapped, f"add these to _OPERATION_BY_CALL_TYPE: {unmapped}"
def test_unmapped_call_type_falls_back_to_chat_loudly(caplog):
"""The fallback still labels the series ``chat`` so it is never unlabelled,
but it says so at debug: a silent default is how retrieval and agent calls
ended up in the chat charts in the first place."""
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
assert resolve_operation("some_future_call_type") is GenAIOperation.CHAT
assert any("some_future_call_type" in record.getMessage() for record in caplog.records)
# --- MCP tool-call (source of truth #1/#2/#3) ------------------------------- #

58
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -4835,66 +4835,26 @@ async def test_common_checks_personal_user_budget_blocks_in_gather():
@pytest.mark.asyncio
async def test_user_budget_enforced_on_team_key():
"""User budget must be enforced even when the key belongs to a team.
async def test_common_checks_personal_user_budget_skipped_for_team_key():
"""A user's personal max_budget does not apply to a team-scoped key.
Previously _user_max_budget_check skipped enforcement for team keys,
letting a user with a $100 personal budget spend unlimited through a
team key. This regression test ensures that is no longer the case.
Team keys are governed by the team (and team-member) budgets only; the key
owner's personal budget is deliberately out of scope. This asserts the read
path lets a team key through even when the user is far over their personal
budget, and fails if personal enforcement is reintroduced for team keys.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0)
team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0)
token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == "spend:user:u1" else 0.0
async def _no_membership(*a, **kw):
return None
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership):
with pytest.raises(litellm.BudgetExceededError) as over:
await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=team,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
assert "User=u1" in str(over.value)
@pytest.mark.asyncio
async def test_skip_user_budget_on_team_key_flag_restores_old_behavior():
"""Setting skip_user_budget_on_team_key=True skips user budget for team keys.
This is the opt-in escape hatch that restores the legacy behavior where
user budgets were not enforced when the key belonged to a team.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0)
token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == "spend:user:u1" else 0.0
async def _no_membership(*a, **kw):
async def _no_membership(*args, **kwargs):
return None
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
@ -4906,7 +4866,7 @@ async def test_skip_user_budget_on_team_key_flag_restores_old_behavior():
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={"skip_user_budget_on_team_key": True},
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),

49
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -611,12 +611,12 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_
@pytest.mark.asyncio
async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_state):
"""A user's personal budget must be reserved even when the key belongs to a team.
async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state):
"""The reservation path mirrors the read path: no personal user counter for a team key.
Regression for GitHub issue #12905: previously the reservation path skipped the
user spend counter whenever the key had a team, so a team key could overshoot the
user's personal max_budget under concurrency.
A team-scoped key reserves against the key and team counters only, so the key
owner's personal max_budget never gates a team request. Fails if the user
counter is reserved for team keys again.
"""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
@ -645,44 +645,7 @@ async def test_should_reserve_user_budget_counter_for_team_key(spend_counter_sta
proxy_logging_obj=proxy_logging_obj,
)
assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") == pytest.approx(0.3)
await release_budget_reservation(reservation)
@pytest.mark.asyncio
async def test_should_skip_user_budget_counter_for_team_key_when_flag_set(spend_counter_state):
"""skip_user_budget_on_team_key=True restores the legacy behavior where a user's
personal budget is not reserved for a team key."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-user-on-team-skip",
spend=0.0,
user_id="user-on-team-skip",
team_id="team-no-budget-skip",
)
team_object = LiteLLM_TeamTable(team_id="team-no-budget-skip", spend=0.0, max_budget=None)
user_object = LiteLLM_UserTable(user_id="user-on-team-skip", spend=0.0, max_budget=5.0)
with patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=0.3,
):
reservation = await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=valid_token,
team_object=team_object,
user_object=user_object,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
skip_user_budget_on_team_key=True,
)
assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-skip") is None
assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") is None
await release_budget_reservation(reservation)

32
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -9201,38 +9201,6 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
app.dependency_overrides.clear()
def test_get_config_list_includes_skip_user_budget_on_team_key(monkeypatch):
"""Related to #12905: the opt-out flag must be discoverable via /config/list so
it renders as a Boolean toggle on the Admin UI General Settings table. This
requires both the ConfigGeneralSettings field and the allowed_args entry."""
import types
from unittest.mock import AsyncMock, MagicMock
from fastapi.testclient import TestClient
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.proxy_server import app
mock_prisma = MagicMock()
mock_config_table = MagicMock()
mock_config_table.find_first = AsyncMock(return_value=None)
mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
client = TestClient(app)
resp = client.get("/config/list", params={"config_type": "general_settings"})
assert resp.status_code == 200, resp.text
fields = {item["field_name"]: item for item in resp.json()}
assert "skip_user_budget_on_team_key" in fields
assert fields["skip_user_budget_on_team_key"]["field_type"] == "Boolean"
finally:
app.dependency_overrides.clear()
def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch):
"""The throttle fraction is a litellm_settings scalar surfaced on the General
Settings table as a Float field so it sits with the other global limits; it

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

@ -22820,11 +22820,6 @@ export interface components {
* @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.
*/
reject_clientside_metadata_tags?: boolean | null;
/**
* Skip User Budget On Team Key
* @description If True, restores the legacy behavior where a user's personal max_budget is NOT enforced when their key belongs to a team; only the team (and team-member) budgets apply. Defaults to False, meaning the user's personal max_budget is always enforced regardless of whether the key belongs to a team (see GitHub issue #12905).
*/
skip_user_budget_on_team_key?: boolean | null;
/**
* Store Model In Db
* @description If True, models and config are stored in and loaded from the database. Default is False.