From 6f1625d23bb1649394ad59c6a98c31633a1ba924 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 30 Jul 2026 13:19:49 -0700 Subject: [PATCH 1/2] revert(proxy)!: stop enforcing user budget on team keys (#35271) Reverts #32005. Team-scoped keys are governed by the team and team-member budgets only; the key owner personal max_budget no longer applies to them, restoring the hierarchy that existed before that PR. The skip_user_budget_on_team_key opt-out existed solely to turn the new behavior back off, so it is removed along with the behavior: the ConfigGeneralSettings field, the /config/list allowed_args entry that surfaced it as an Admin UI toggle, and the argument threaded through reserve_budget_for_request and _get_budget_counters. Regression tests cover both enforcement points in the restored direction: test_common_checks_personal_user_budget_skipped_for_team_key for the read-time check and test_should_not_reserve_user_budget_counter_for_team_key for the optimistic reservation path. --- litellm/proxy/_types.py | 10 ---- litellm/proxy/auth/auth_checks.py | 41 ++++++------- litellm/proxy/auth/user_api_key_auth.py | 1 - litellm/proxy/proxy_server.py | 1 - .../spend_tracking/budget_reservation.py | 6 +- .../test_user_api_key_auth.py | 11 ++-- .../proxy/auth/test_auth_checks.py | 58 +++---------------- .../proxy/test_budget_reservation.py | 49 ++-------------- tests/test_litellm/proxy/test_proxy_server.py | 32 ---------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 -- 10 files changed, 42 insertions(+), 172 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dc84a4ca705..9e98cb46b9a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 03e5e80288e..0472b496b78 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b268d2f8840..f4d07c1a674 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 63748705f0d..c72e3d4ee5b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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"}, } diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 013873179c6..12ecd2ec300 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 59c4caefa33..01dbb65a648 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 34a353966bf..4d0ef58b7f8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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(), diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index e6ccbc579b5..1a584423fac 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -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) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 62f5ced7a39..5646d202e31 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bc3110d4d3d..380b6545da8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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. From abd239f9035b587df77d05dc8043153b8c95f581 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 30 Jul 2026 13:48:59 -0700 Subject: [PATCH 2/2] fix(otel): label retrieval and agent metrics correctly and emit gen_ai.provider.name (#35151) * fix(otel): label retrieval and agent metrics correctly and emit gen_ai.provider.name The GenAI metric attribute builder mapped only chat, text completion, embedding, responses and MCP tool calls to an operation name, so vector-store searches and A2A agent sends fell through to the "chat" default. Their duration and cost then landed in the same series a Grafana GenAI dashboard reads chat latency off, with no way to tell them apart. Both now map to the operation names the convention defines for them, retrieval and invoke_agent, and an unmapped call type says so at debug instead of silently becoming chat. The provider label used gen_ai.system, which the convention deprecated in favor of gen_ai.provider.name; the dashboards built on that vocabulary find nothing under the old key. Metrics now carry gen_ai.provider.name with the semconv provider value (bedrock -> aws.bedrock) via the resolve_provider helper the span path already uses, and keep dual-emitting gen_ai.system with its raw value so a dashboard already querying it keeps matching. A request litellm cannot attribute to a provider gets no provider label at all rather than a placeholder "Unknown" that minted a permanent series nobody can act on. Resolves LIT-4954 Resolves LIT-4959 * fix(otel): map the rest of the vector-store call types off the chat default Mapping only the search left the store lifecycle (create, retrieve, list, update, delete) and the file operations (create, list, retrieve, content, update, delete) falling through to chat, so vector-store admin traffic kept polluting the same series a dashboard reads chat latency off. A live run confirmed it: all 20 metric datapoints from a create, retrieve, list, file-list and delete came out labelled chat. The convention names no operation for vector-store management, so these take vendor values under the litellm. prefix, litellm.vector_store_management and litellm.vector_store_file_management, one per REST resource. Its note on gen_ai.operation.name directs instrumentation to use a system-specific name when no predefined value applies, which is the same allowance resolve_provider already relies on for unmapped providers. Excluding them from the GenAI metrics altogether was the alternative; it deletes series an operator may be watching today and is far harder to reverse than a rename, so it stays available as a follow-up rather than being decided here. Mapping them onto the semconv memory store family was rejected: litellm vector stores hold documents, not agent memory records, and borrowing those names would put document admin calls into whatever charts agent-memory operations, which is the bug this fixes. /rag/query reaches the same recorder and is the same operation as a vector-store search, so query and aquery map to retrieval too; leaving them would have left the defect alive on a second retrieval surface. /rag/ingest is a write with no semconv equivalent and no retrieval or agent confusion, so it is left for the RAG owners to name. Resolves LIT-4954 * fix(otel): give the streaming A2A path a call type so it labels as invoke_agent The streaming logging object is built by hand and never runs through update_environment_variables, the only place call_type reaches model_call_details, so every streamed agent turn arrived at the recorder with no call type and fell back to chat. Stamp it, and map the streaming spelling alongside the non-streaming ones. --- litellm/a2a_protocol/main.py | 1 + litellm/integrations/opentelemetry.py | 1 + litellm/integrations/otel/model/semconv.py | 72 ++++++++++- litellm/integrations/otel/plumbing/metrics.py | 33 ++++- tests/test_litellm/a2a_protocol/test_main.py | 29 +++++ .../integrations/otel/test_otel_v2_metrics.py | 117 +++++++++++++++++- .../otel/test_otel_v2_sources_of_truth.py | 102 +++++++++++++++ 7 files changed, 339 insertions(+), 16 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 4c23ecfed54..f04edf2579b 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -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 diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 11e7ab062b6..60eb960be4e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -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", diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index a0994d1948a..3b38b3d3655 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -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 diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 4d3c39e33d4..6ebaacafdc4 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -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", } diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 2a675616245..c5e48620fde 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -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" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py index 56607414de4..e1b8e4b5721 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py @@ -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", diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 71be28ea485..612ac1e5113 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -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) ------------------------------- #