diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml index af23512dafc..203082f27ba 100644 --- a/helm/litellm-helm/tests/connection_pool_tests.yaml +++ b/helm/litellm-helm/tests/connection_pool_tests.yaml @@ -59,3 +59,54 @@ tests: - contains: path: spec.template.spec.containers[0].args content: "4" + + - it: should give the collector sidecar the same pool env as the proxy container + template: deployment.yaml + set: + collector.enabled: true + db.connectionPool.enabled: true + db.connectionPool.maxDbConnections: 8 + db.connectionPool.maxClientConn: 400 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "400" + + - it: should give the collector sidecar no pool env when the pool is off + template: deployment.yaml + set: + collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 5c1a089b50a..c06cc9583a0 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -170,6 +170,9 @@ spec: - name: CONFIG_FILE_PATH value: /app/config/config.yaml {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} - name: LITELLM_JOB_ROLE value: collector diff --git a/helm/litellm/tests/collector_tests.yaml b/helm/litellm/tests/collector_tests.yaml index b1199a16793..4ef7e3c8ca4 100644 --- a/helm/litellm/tests/collector_tests.yaml +++ b/helm/litellm/tests/collector_tests.yaml @@ -29,16 +29,16 @@ tests: value: Resource template: gateway/hpa.yaml - - it: runs the collector as a sidecar sharing env, config and a unix socket emptyDir, and scales on the gateway container only + - it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only set: gateway.collector.enabled: true gateway.collector.bufferSize: 250 gateway.collector.onUnavailable: drop gateway.image.tag: v1.102.0 gateway.numWorkers: 4 - gateway.extraEnv: - - name: LITELLM_PGBOUNCER_ENABLED - value: "true" + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 gateway.envSecrets: - litellm-license gateway.volumes: @@ -107,6 +107,18 @@ tests: name: LITELLM_PGBOUNCER_ENABLED value: "true" template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + template: gateway/deployment.yaml - contains: path: spec.template.spec.containers[1].env content: diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml index 6dd9c274a8c..c39651a52c9 100644 --- a/helm/litellm/tests/connection_pool_tests.yaml +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -82,9 +82,70 @@ tests: name: LITELLM_PGBOUNCER_ENABLED any: true + - it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + gateway.metricsServer.enabled: true + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: metrics + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - equal: + path: spec.template.spec.containers[2].name + value: collector + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + + - it: collector sidecar gets no pool env when the pool is off + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true + - it: pool with IAM auth renders both the pool and the token auth flag template: gateway/deployment.yaml set: + gateway.collector.enabled: true database.connectionPool.enabled: true database.writer.useIAMAuth: true asserts: @@ -98,6 +159,16 @@ tests: content: name: IAM_TOKEN_DB_AUTH value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" - it: pool with Entra auth renders both the pool and the token auth flag template: gateway/deployment.yaml diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 1f67b984513..1873219d1ea 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -234,8 +234,8 @@ database: # workers run; the workers connect to the pool over loopback, with no extra # network hop. The chart emits LITELLM_PGBOUNCER_ENABLED / # LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on - # the gateway container only: the backend runs a single worker and the - # migrations Job must keep a direct connection. With + # the gateway container and its collector sidecar only: the backend runs a + # single worker and the migrations Job must keep a direct connection. With # `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and # renews the database token itself, so the workers never see it. Starting profile for # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 91b4e4a7ba1..7d4c78088f1 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.95" +version = "0.4.96" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.95" +version = "0.4.96" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 893bfeefff8..6ee68ab21c5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -295,6 +295,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS +_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( + ("user_api_key_auth", "user_api_key_budget_reservation") +) sentry_sdk_instance = None capture_exception = None @@ -5386,23 +5389,23 @@ class StandardLoggingPayloadSetup: Returns: dict: Merged metadata with user API key fields taking precedence """ - merged_metadata: Final[dict] = {} - - # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict): - for key, value in litellm_params["metadata"].items(): - # Skip non-serializable objects like UserAPIKeyAuth - if key in {"user_api_key_auth", "user_api_key_budget_reservation"}: - continue - merged_metadata[key] = value - - # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict): - for key, value in litellm_params["litellm_metadata"].items(): - if key not in merged_metadata: # Don't overwrite existing keys from metadata - merged_metadata[key] = value - - return merged_metadata + metadata: Final = litellm_params.get("metadata") + litellm_metadata: Final = litellm_params.get("litellm_metadata") + user_metadata: Final = MappingProxyType( + { + key: value + for key, value in (metadata.copy().items() if isinstance(metadata, dict) else ()) + if key not in _UNSERIALIZABLE_METADATA_KEYS + } + ) + model_metadata: Final = MappingProxyType( + { + key: value + for key, value in (litellm_metadata.copy().items() if isinstance(litellm_metadata, dict) else ()) + if key not in user_metadata + } + ) + return {**user_metadata, **model_metadata} # mutable-ok: function contract returns a plain dict @staticmethod def get_standard_logging_metadata( @@ -5660,7 +5663,7 @@ class StandardLoggingPayloadSetup: additional_logging_headers[key] = additiona_headers[_key] # Preserve all remaining headers verbatim (e.g. llm_provider-x-request-id) - for k, v in additiona_headers.items(): + for k, v in additiona_headers.copy().items(): if k.lower() not in typed_keys: additional_logging_headers[k] = v diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index db23929e0c3..6a5a8832cc6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -236,9 +236,11 @@ class CustomStreamWrapper: stream_options=None, make_call: Callable | None = None, _response_headers: dict | httpx.Headers | None = None, + count_prompt_tokens: Callable[[], int] | None = None, ): self.model = model self.make_call = make_call + self.count_prompt_tokens = count_prompt_tokens self.custom_llm_provider = custom_llm_provider self.logging_obj: LiteLLMLoggingObject = logging_obj self.completion_stream = completion_stream @@ -1641,7 +1643,7 @@ class CustomStreamWrapper: except Exception: model_response.choices[0].delta = Delta() else: - if self.stream_options is not None and self.stream_options["include_usage"] is True: + if self.send_stream_usage is True: model_response.choices = [] return model_response self._record_usage_only_chunk(model_response=model_response) @@ -1996,6 +1998,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) except Exception as e: # stream_chunk_builder can re-raise (as APIError) on large agentic @@ -2248,6 +2251,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) except Exception as e: # see sync __next__: a raise from stream_chunk_builder inside this @@ -2371,6 +2375,7 @@ class CustomStreamWrapper: chunks=self.chunks, messages=self.messages if isinstance(self.messages, list) else None, logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, ) if partial_response is None: return diff --git a/litellm/main.py b/litellm/main.py index 1a4beb787bc..17edafcdfca 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -850,6 +850,12 @@ def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None: ) +def admitted_prompt_token_counter(prompt_tokens: int | None) -> Callable[[], int] | None: + if prompt_tokens is None: + return None + return lambda: prompt_tokens + + def mock_completion( model: str, messages: list, @@ -935,23 +941,26 @@ def mock_completion( if stream is True: model_response = ModelResponseStream() + count_prompt_tokens: Final = admitted_prompt_token_counter(prompt_tokens) # don't try to access stream object, if kwargs.get("acompletion", False) is True: return CustomStreamWrapper( completion_stream=async_mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", logging_obj=logging, + count_prompt_tokens=count_prompt_tokens, ) return CustomStreamWrapper( completion_stream=mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", logging_obj=logging, + count_prompt_tokens=count_prompt_tokens, ) if isinstance(mock_response, litellm.MockException): raise mock_response diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index a074f02f4e8..0339cf4dfea 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -59,6 +59,10 @@ def _get_priority_settings() -> "PriorityReservationSettings": return settings +def _is_latin1_encodable(value: object) -> bool: + return all(ord(char) < 256 for char in str(value)) + + class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): """ Saturation-aware priority-based rate limiter using v3 infrastructure. @@ -666,7 +670,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if response_has_hidden_params(response): priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) additional_headers: Final = ensure_response_additional_headers(response) - additional_headers["x-litellm-priority"] = priority or "default" + priority_header: Final = priority or "default" + if _is_latin1_encodable(priority_header): + additional_headers["x-litellm-priority"] = priority_header + else: + verbose_proxy_logger.debug( + "Skipping x-litellm-priority header: priority %r is not Latin-1 encodable", priority + ) additional_headers["x-litellm-rate-limiter-version"] = "v3" return response diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 38dfd143cbc..1a4764c291e 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -455,6 +455,13 @@ If 2+ reasoning markers are detected in the user message, the request is promote Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier. +For requests identified by a `claude-cli/` or `claude-code/` user agent, the LLM classifier omits caller system +text to avoid classifying environment, agent, and skill catalogs. The current ask, configured prior-turn context, +and trajectory signal remain unchanged. The routed completion still receives the original system text. This +also excludes genuine task constraints supplied only in Claude Code system messages. Other clients keep the +existing system-context behavior. The browser routing preview has no client-identity field and retains that +generic behavior; use the real client when checking Claude Code routing. + ### Harness Reminder Blocks Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a748bb5df27..7519f4c5156 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -48,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( request_contains_image_content, ) from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload +from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.tier_predictor import ( @@ -1993,8 +1994,7 @@ class ComplexityRouter(CustomLogger): Args: prompt: The current user ask text (already extracted as the real human ask, not tool results) - system_prompt: The caller's system prompt (task constraints), always included so later - turns never lose it + system_prompt: Caller task constraints, omitted from classification for Claude Code requests request_kwargs: Request metadata for spend attribution messages: Full message history for extracting prior turns and the trajectory signal """ @@ -2027,9 +2027,18 @@ class ComplexityRouter(CustomLogger): ) encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs) + caller_system_prompt: Final = ( + None + if any( + is_claude_code_user_agent(user_agent) + for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ()) + if isinstance(user_agent := metadata.get("user_agent"), str) + ) + else system_prompt + ) user_payload: Final = self._build_classifier_user_payload( prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, - system_prompt=system_prompt, + system_prompt=caller_system_prompt, prior_turns=prior_turns, messages=messages, has_prior_conversation=has_prior_conversation, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fb7003c887c..1f1b5a5cc4b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -971,9 +971,11 @@ class ComplexityRouterConfig(BaseModel): "classified against what it refers to. Counts turns of both roles when " "classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier " "model, which may " - "be a different deployment or provider than the routed completion model; that call already " - "carries the current user ask and the caller's system prompt in full. Set to 0 to send neither " - "prior turns nor any conversation context beyond the current ask. Only applies when " + "be a different deployment or provider than the routed completion model; that call carries " + "the current user ask and, except for Claude Code requests, the extracted system-role text in full. " + "Claude Code system text is omitted to avoid classifying harness instructions; the routed " + "completion still receives it. Set to 0 to send neither prior turns nor " + "any conversation context beyond the current ask. Only applies when " "classifier_type is 'llm'." ), ) @@ -985,9 +987,9 @@ class ComplexityRouterConfig(BaseModel): "context window, per classification call. Turns are taken newest first and quoted whole " "while they fit, so a conversation small enough to quote entirely is never cut; once the " "budget runs out the older turns are dropped whole and only the turn straddling the " - "boundary is truncated, into whatever space is left. The current ask and the caller's " - "system prompt sit outside this budget and are always sent in full, as does the numbering " - "each quoted turn carries. A budget under 120 leaves no room to quote a turn and " + "boundary is truncated, into whatever space is left. The current ask and, except for Claude " + "Code requests, the extracted system-role text sit outside this budget and are sent in full, as does " + "the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and " "suppresses the block; set classifier_context_window_size to 0 to turn context off " "deliberately. Only applies when classifier_type is 'llm'." ), diff --git a/litellm/utils.py b/litellm/utils.py index aced4c9b312..a765e1b1246 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -66,6 +66,7 @@ from litellm.constants import ( DEFAULT_EMBEDDING_PARAM_VALUES, DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, @@ -278,7 +279,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args from litellm import utils as litellm_utils @@ -1195,6 +1196,47 @@ def function_setup( raise e +def _dispatch_success_logging( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + is_completion_with_fallbacks: bool, + is_litellm_internal_call: bool, +) -> None: + if not is_litellm_internal_call: + if getattr(logging_obj, "_defer_async_logging", False): + + def _enqueue_deferred_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging + else: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=result, + start_time=start_time, + end_time=end_time, + ) + + async def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, result, @@ -1662,6 +1704,16 @@ def client(original_function): kwargs=kwargs, ) + _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") # Copy the current context to propagate it to the background thread @@ -1676,15 +1728,6 @@ def client(original_function): end_time, ) # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") - update_response_metadata( - result=result, - logging_obj=logging_obj, - model=model, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) return result except Exception as e: call_type = original_function.__name__ @@ -1944,48 +1987,20 @@ def client(original_function): args=args, ) - # LOG SUCCESS - handle streaming success logging in the _next_ object - # Internal sub-calls (e.g. emulated file-search steps) share the - # parent's logging obj; skip async logging here so only the outer call bills once. - # NOTE: streaming requests return early (before this point) via - # CustomStreamWrapper, so this block is non-streaming only. - if not _is_litellm_internal_call: - if getattr(logging_obj, "_defer_async_logging", False): - - def _enqueue_deferred_logging() -> None: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging - else: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=result, - start_time=start_time, - end_time=end_time, - ) # REBUILD EMBEDDING CACHING if ( isinstance(result, EmbeddingResponse) and _caching_handler_response is not None and _caching_handler_response.final_embedding_cached_response is not None ): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return _llm_caching_handler._combine_cached_embedding_response_with_api_result( _caching_handler_response=_caching_handler_response, embedding_response=result, @@ -2001,6 +2016,14 @@ def client(original_function): start_time=start_time, end_time=end_time, ) + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return result except Exception as e: @@ -7001,7 +7024,26 @@ class TextCompletionStreamWrapper: raise StopAsyncIteration -def mock_completion_streaming_obj(model_response, mock_response, model, n: int | None = None): +def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream: + return ModelResponseStream( + id=model_response.id, + choices=[], # mutable-ok: ModelResponseStream only treats a list as explicit choices, a tuple gets a default choice + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + ), + ) + + +def mock_completion_streaming_obj( + model_response: ModelResponseStream, + mock_response: str | MockException | ModelResponseStream, + model: str, + n: int | None = None, + prompt_tokens: int | None = None, +) -> Iterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7021,14 +7063,17 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) async def async_mock_completion_streaming_obj( - model_response, + model_response: ModelResponseStream, mock_response: str | MockException | ModelResponseStream, - model, + model: str, n: int | None = None, -): + prompt_tokens: int | None = None, +) -> AsyncIterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7048,6 +7093,8 @@ async def async_mock_completion_streaming_obj( _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) ########## Reading Config File ############################ diff --git a/pyproject.toml b/pyproject.toml index 04f2f3fd1dd..d33d693f794 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.95", + "litellm-proxy-extras==0.4.96", "litellm-enterprise==0.1.66", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 67f1270d0c2..6fcbdc2f500 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -362,7 +362,14 @@ The sidecar gets the same database, Redis, master-key, license, proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values as the gateway container, runs with `LITELLM_JOB_ROLE=collector`, and is non-essential with an ECS restart policy, so a sidecar crash restarts it in -place while the gateway falls back to in-process spend tracking. +place while the gateway falls back to in-process spend tracking. With +`gateway_connection_pool_enabled` it also gets the `LITELLM_PGBOUNCER_*` env, +so with a password-authenticated database (`create_database = false`) its +Prisma client goes through the task-local PgBouncer instead of opening a +second pool straight to the database. Under IAM token auth (the module-managed +Aurora cluster) the collector keeps its own direct connection on purpose: the +pooler's auth file only holds the token the gateway container minted, which +the sidecar cannot present, so it mints its own. ```hcl collector_enabled = true diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl index ab0a1cfc076..7408b1179ed 100644 --- a/terraform/litellm/aws/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -51,6 +51,43 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" { } } +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for env in [local.gateway_environment, local.collector_container[0].environment] : ( + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) + ]) + error_message = "The collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the task-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue([for e in local.collector_container[0].environment : startswith(e.name, "LITELLM_PGBOUNCER_")]) + error_message = "The collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + run "pool_enabled_uses_the_module_default_sizes" { command = plan diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 23cfe9afc24..4b2f576adc1 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -339,8 +339,11 @@ instead of the Unix socket helm uses; the proxy rejects any non-loopback address. The sidecar runs the same Redis CA + `DATABASE_URL` bootstrap as the gateway container, gets the same database, Redis, master-key, license, proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values, and -runs with `LITELLM_JOB_ROLE=collector`. When it is unreachable the -gateway falls back to in-process spend tracking. +runs with `LITELLM_JOB_ROLE=collector`. With `gateway_connection_pool_enabled` +it also gets the `LITELLM_PGBOUNCER_*` env, so its Prisma client goes through +the instance-local PgBouncer instead of opening a second pool straight to the +database. When it is unreachable the gateway falls back to in-process spend +tracking. ```hcl collector_enabled = true diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl index 439fe6b1d71..999e4f0ff95 100644 --- a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -79,6 +79,45 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" { } } +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : ( + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) if c.name == "spend-collector" + ]) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the instance-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue(flatten([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : [ + for e in c.env : startswith(e.name, "LITELLM_PGBOUNCER_") + ] if c.name == "spend-collector" + ])) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + run "pool_enabled_uses_the_module_default_sizes" { command = plan diff --git a/tests/llm_translation/test_azure_o_series.py b/tests/llm_translation/test_azure_o_series.py index 1a2d672af71..ce7e614cbe2 100644 --- a/tests/llm_translation/test_azure_o_series.py +++ b/tests/llm_translation/test_azure_o_series.py @@ -159,15 +159,23 @@ def test_azure_o_series_routing(): def test_openai_o_series_max_retries_0(mock_get_openai_client): import litellm + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( + ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + ) litellm.set_verbose = True response = litellm.completion( model="azure/o1-preview", messages=[{"role": "user", "content": "hi"}], max_retries=0, + api_key="fake-key", + api_base="https://fake-azure.openai.azure.com", + api_version="2024-10-21", ) mock_get_openai_client.assert_called_once() assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0 + assert response.choices[0].message.content == "Hello" @pytest.mark.asyncio diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 0fa72b45ed8..e6528e77749 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -335,6 +335,10 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): ] with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + mock_post.return_value.headers = {} + mock_post.return_value.parse.return_value = litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}] + ) response = litellm.completion( model="azure/gpt-4.1-mini", messages=[ @@ -362,6 +366,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): assert "response_format" in mock_post.call_args.kwargs else: assert "response_format" not in mock_post.call_args.kwargs + assert response.choices[0].message.content == InvestigationOutput().model_dump_json() def test_map_openai_params(): diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 2b9abdec5d0..af4ba85d58e 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -292,15 +292,21 @@ class TestOpenAIChatCompletion(BaseLLMChatTest): def test_openai_max_retries_0(mock_get_openai_client): import litellm + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} + mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( + ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + ) litellm.set_verbose = True response = litellm.completion( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], max_retries=0, + api_key="fake-key", ) mock_get_openai_client.assert_called_once() assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0 + assert response.choices[0].message.content == "Hello" @patch("litellm.main.openai_chat_completions._get_openai_client") diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index f8f23ea015a..43ed57f63af 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -3999,10 +3999,14 @@ def test_completion_novita_ai(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client.chat.completions.with_raw_response, "create" ) as mock_call: + mock_call.return_value.headers = {} + mock_call.return_value.parse.return_value = litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": "Hello"}}] + ) try: - completion( + response = completion( model="novita/meta-llama/llama-3.3-70b-instruct", messages=messages, client=openai_client, @@ -4010,6 +4014,7 @@ def test_completion_novita_ai(): ) mock_call.assert_called_once() + assert response.choices[0].message.content == "Hello" # Verify model is passed correctly assert ( diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index f0f24a6e6b2..834570091bd 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -1076,7 +1076,7 @@ def test_standard_logging_payload(model, turn_off_message_logging): ) ) - keys_list = list(StandardLoggingPayload.__annotations__.keys()) + keys_list = list(StandardLoggingPayload.__required_keys__) for k in keys_list: assert ( @@ -1190,7 +1190,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream): ) ) - keys_list = list(StandardLoggingPayload.__annotations__.keys()) + keys_list = list(StandardLoggingPayload.__required_keys__) for k in keys_list: assert ( diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index 83a652e8884..7ac9ac0b5ad 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -270,7 +270,7 @@ async def test_datadog_logging_http_request(): message = json.loads(body[0]["message"]) print("logged message", json.dumps(message, indent=4)) - expected_message_fields = StandardLoggingPayload.__annotations__.keys() + expected_message_fields = StandardLoggingPayload.__required_keys__ for field in expected_message_fields: assert field in message, f"Field '{field}' is missing from the message" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index aeaa8f18c26..2f6339dcdbb 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3,6 +3,7 @@ import contextlib import datetime import os import sys +from collections.abc import Callable from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -6945,3 +6946,61 @@ def test_classifier_audit_is_not_added_to_other_calls(logging_obj, call_type, or logging_obj.model_call_details["litellm_params"] = {"metadata": {"internal_call_origin": origin}} logging_obj.pre_call(input=[], api_key=None, additional_args={"complete_input_dict": {"input": "embedding"}}) assert logging_obj.classifier_input is None + + +def _run_while_a_thread_grows(target: dict, read: Callable[[], None], reads: int) -> None: + import itertools + import threading + + stop: Final = threading.Event() + + def grow() -> None: + for counter in itertools.count(): + if stop.is_set(): + return + key: Final = f"late_{counter % 64}" + if key in target: + del target[key] + else: + target[key] = counter + + writer: Final = threading.Thread(target=grow, daemon=True) + previous_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + writer.start() + try: + for _ in range(reads): + read() + finally: + stop.set() + writer.join(timeout=5) + sys.setswitchinterval(previous_interval) + + +def test_merge_litellm_metadata_survives_a_thread_growing_metadata_mid_merge(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + metadata: Final = {f"key_{i}": i for i in range(2000)} + litellm_params: Final = {"metadata": metadata, "litellm_metadata": {"model_group": "gpt"}} + + def read() -> None: + merged: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) + assert merged["key_1999"] == 1999 + assert merged["model_group"] == "gpt" + + _run_while_a_thread_grows(metadata, read, reads=300) + + +def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + headers: Final = {f"llm_provider-x-custom-{i}": str(i) for i in range(2000)} + headers["x-ratelimit-remaining-requests"] = "7" + + def read() -> None: + copied: Final = StandardLoggingPayloadSetup.get_additional_headers(headers) + assert copied is not None + assert copied["x_ratelimit_remaining_requests"] == 7 + assert copied["llm_provider-x-custom-1999"] == "1999" + + _run_while_a_thread_grows(headers, read, reads=300) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 626b8a63b20..efe4209c1c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1554,3 +1554,52 @@ def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None assert response is not None assert response.choices[0].message.role == "user" assert response.choices[0].message.content == "Hi" + + +def _fail_prompt_token_count() -> int: + raise AssertionError("prompt tokens must come from the usage chunk, not the tokenizer") + + +def test_calculate_usage_reads_prompt_tokens_from_mock_stream_usage_chunk_without_tokenizer_fallback() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), + mock_response="ok", + model="gpt-5.4-mini", + prompt_tokens=51234, + ) + ) + assert chunks[-1].choices == [] + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=_fail_prompt_token_count, + ) + + assert usage.prompt_tokens == 51234 + assert usage.completion_tokens == chunks[-1].usage.completion_tokens + assert usage.total_tokens == 51234 + usage.completion_tokens + + +def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_admission_count() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), mock_response="ok", model="gpt-5.4-mini" + ) + ) + assert all(chunk.choices for chunk in chunks) + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=lambda: 77, + ) + + assert usage.prompt_tokens == 77 diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0cd6b4ede9c..527449bbc48 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1918,3 +1918,30 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): ) assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_metadata", "expected_priority_header"), + [ + ({"priority": "优先"}, None), + ({"priority": "high"}, "high"), + ({}, "default"), + ], +) +async def test_post_call_success_hook_priority_header_is_always_http_encodable(team_metadata, expected_priority_header): + from starlette.responses import Response + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + response = {"id": "msg_123", "type": "message", "role": "assistant", "content": [], "_hidden_params": {}} + + await handler.async_post_call_success_hook( + data={"model": "anthropic-haiku"}, + user_api_key_dict=UserAPIKeyAuth(team_id="team-1", team_metadata=team_metadata), + response=response, + ) + + additional_headers = response["_hidden_params"]["additional_headers"] + http_response = Response(headers={key: str(value) for key, value in additional_headers.items()}) + assert http_response.headers.get("x-litellm-priority") == expected_priority_header + assert http_response.headers["x-litellm-rate-limiter-version"] == "v3" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index b8e728c44a9..dddaee71a63 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2676,6 +2676,26 @@ class TestEncryptedTaskClassifier: assert "source-secret" not in json.dumps(call) assert "originating_request_masked" not in call["proxy_server_request"]["body"] + @pytest.mark.asyncio + async def test_claude_code_encrypted_task_omits_caller_instructions(self): + router, dependency = _native_classifier_router() + task: Final = _encrypted_agent_task() + request: Final = { + "input": [task], + "instructions": "CLAUDE_CODE_SYSTEM", + "litellm_metadata": {"user_agent": "claude-cli/2.1.233"}, + } + original: Final = deepcopy(request) + + result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request) + + assert result.routing_decision["cause"] == "llm_classifier" + assert request == original + call: Final = dependency.aresponses.call_args.kwargs + assert call["instructions"] == classification_system_prompt(router.config.classifier_context_window_size) + assert "CLAUDE_CODE_SYSTEM" not in json.dumps(call["input"][:-1]) + assert call["input"][-1] == task + @pytest.mark.asyncio @pytest.mark.parametrize( "items", @@ -7994,6 +8014,113 @@ _CODEX_ENVELOPES: Final = ( class TestContextAwareClassifier: """Test the new classifier context window and trajectory signals.""" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_metadata,forwards_system", + [ + ({"metadata": {"user_agent": "claude-cli/2.1.233"}}, False), + ({"litellm_metadata": {"user_agent": "claude-code/2.1.233"}}, False), + ({"metadata": {"user_agent": "curl/8.7.1"}}, True), + ({"litellm_metadata": {}}, True), + ( + {"metadata": {"user_agent": "claude-cli/2.1.233"}, "litellm_metadata": {"user_agent": "curl/8.7.1"}}, + False, + ), + ({"metadata": {"user_agent": "Claude-Code/2.1.233"}}, True), + ], + ) + async def test_claude_code_classifier_omits_harness_system_prompt( + self, + llm_classifier_config: dict[str, object], + request_metadata: dict[str, object], + forwards_system: bool, + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + { + **llm_classifier_config, + "classifier_context_include_assistant_turns": True, + }, + ) + messages: Final = [ + {"role": "user", "content": "Design the retry state machine"}, + {"role": "assistant", "content": "The design needs a lease and fencing token"}, + {"role": "user", "content": "Now prove it cannot livelock"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + top_level_system: Final = [{"type": "text", "text": "TOP_LEVEL_HARNESS_SYSTEM"}] + claude_kwargs: Final = { + "metadata": {"user_agent": "claude-cli/2.1.233"}, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + compared_kwargs: Final = { + **request_metadata, + "system": top_level_system, + "proxy_server_request": {"body": {"system": top_level_system}}, + } + original_messages: Final = deepcopy(messages) + original_kwargs: Final = deepcopy((claude_kwargs, compared_kwargs)) + results: Final = ( + await router.async_pre_routing_hook("test-complexity-router", claude_kwargs, messages), + await router.async_pre_routing_hook("test-complexity-router", compared_kwargs, messages), + ) + + assert all(result is not None and result.routing_decision["cause"] == "llm_classifier" for result in results) + assert all(result is not None and result.messages == original_messages for result in results) + assert messages == original_messages + assert (claude_kwargs, compared_kwargs) == original_kwargs + calls: Final = tuple(call.kwargs["messages"] for call in dependency.acompletion.await_args_list) + assert calls[0][0]["content"] == calls[1][0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + payloads: Final = (calls[0][1]["content"], calls[1][1]["content"]) + for payload, expected_system in zip(payloads, (False, forwards_system)): + assert payload.endswith("Classify this message:\nNow prove it cannot livelock") + assert ("ENVIRONMENT_CATALOG" in payload) is expected_system + assert ("AGENT_CATALOG" in payload) is expected_system + assert ("SKILL_CATALOG" in payload) is expected_system + assert "Design the retry state machine" in payload + assert "lease and fencing token" in payload + assert "TOP_LEVEL_HARNESS_SYSTEM" not in payload + assert "Conversation so far: ~35 tokens across the request" in payload + + @pytest.mark.asyncio + async def test_claude_code_first_turn_without_context_omits_harness_system_prompt( + self, llm_classifier_config: dict[str, object] + ) -> None: + dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))) + router: Final = ComplexityRouter( + "test-complexity-router", + dependency, + {**llm_classifier_config, "classifier_context_window_size": 0}, + ) + messages: Final = [ + {"role": "user", "content": "What is two plus two?"}, + { + "role": "system", + "content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}], + }, + ] + request_kwargs: Final = {"litellm_metadata": {"user_agent": "claude-code/2.1.233"}} + original: Final = deepcopy((messages, request_kwargs)) + + result: Final = await router.async_pre_routing_hook("test-complexity-router", request_kwargs, messages) + + assert result is not None and result.routing_decision["cause"] == "llm_classifier" + assert result.messages == messages == original[0] + assert request_kwargs == original[1] + classifier_messages: Final = dependency.acompletion.call_args.kwargs["messages"] + assert classifier_messages[0]["content"] == classification_system_prompt( + router.config.classifier_context_window_size + ) + assert classifier_messages[1]["content"].strip() == "Classify this message:\nWhat is two plus two?" + @pytest.mark.parametrize( "tail,expected", ( diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index f8841af9750..f71225c6fc5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2421,6 +2421,253 @@ def test_mock_completion_usage_falls_back_to_default_without_admission_count(): assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT +_ADMISSION_INPUT_TOKENS: Final = 51234 + + +def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata + return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}} + + +_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS) +_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] +_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" + + +def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: + return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] + + +def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: + return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + +@pytest.mark.parametrize("n", (None, 2)) +def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", (None, 2)) +async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( + n: int | None, +): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + metadata=_ADMISSION_METADATA, + ) + ) + + assert _client_usage_chunks(chunks) == [] + assert all(len(chunk.choices) == 1 for chunk in chunks) + assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + metadata=_ADMISSION_METADATA, + ) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks: Final = [chunk async for chunk in response] + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok" + assert _client_usage_chunks(chunks) == [] + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + ) + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +def _usage_triple(usage: Usage) -> tuple[int, int, int]: + return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) + + +@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0)) +def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int): + metadata: Final = _admission_metadata(input_tokens) + non_stream: Final = litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + metadata=metadata, + ) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks: Final = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0]) + assert non_stream.usage.prompt_tokens == input_tokens + assert _prompt_token_counter_calls(token_counter) == [] + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response: Final = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=[{"role": "user", "content": ""}], + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + litellm_metadata=_admission_metadata(0), + ) + chunks: Final = [chunk async for chunk in response] + + usage_chunks: Final = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens) + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage(): + metadata: Final = _admission_metadata(0) + non_stream: Final = litellm.text_completion( + model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata + ) + chunks: Final = list( + litellm.text_completion( + model="openai/gpt-5.4-mini", + prompt="", + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=metadata, + ) + ) + + stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) + assert len(stream_usages) == 1 + assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0]) + assert non_stream.usage.prompt_tokens == 0 + + def test_mock_completion_stream_with_model_response(): """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" from litellm import completion diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8b186be43e5..7753cb7d770 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,9 +1,12 @@ import asyncio +import contextlib import json import logging import os +import queue import threading from datetime import datetime, timedelta, timezone +from collections.abc import Iterator from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -15,6 +18,7 @@ from jsonschema import validate import litellm from litellm._internal_context import is_internal_call +from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm._logging import ( CorrelationContextFilter, JsonFormatter, @@ -6208,3 +6212,135 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with "api_key": "sk-from-db", } assert _credential_warnings(caplog) == [] + + +_MOCK_STREAM_ID: Final = "chatcmpl-mock-stream" +_ChunkSnapshot = tuple[str, tuple[str | None, ...], Usage | None] + + +def _snapshot(chunk: ModelResponseStream) -> _ChunkSnapshot: + return chunk.id, tuple(choice.delta.content for choice in chunk.choices), getattr(chunk, "usage", None) + + +def _mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import mock_completion_streaming_obj + + return [ + _snapshot(chunk) + for chunk in mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +async def _async_mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import async_mock_completion_streaming_obj + + return [ + _snapshot(chunk) + async for chunk in async_mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +_CONTENT_SNAPSHOTS: Final = [(_MOCK_STREAM_ID, (content,), None) for content in ("hel", "lo ", "wor", "ld")] + + +def _assert_trailing_usage_chunk(snapshots: list[_ChunkSnapshot], prompt_tokens: int) -> None: + assert snapshots[:-1] == _CONTENT_SNAPSHOTS + chunk_id, choices, usage = snapshots[-1] + assert chunk_id == _MOCK_STREAM_ID + assert choices == () + assert usage is not None + assert usage.prompt_tokens == prompt_tokens + assert usage.completion_tokens == DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage.total_tokens == prompt_tokens + usage.completion_tokens + + +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +def test_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens(prompt_tokens: int) -> None: + _assert_trailing_usage_chunk(_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +async def test_async_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens( + prompt_tokens: int, +) -> None: + _assert_trailing_usage_chunk(await _async_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +def test_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert _mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert await _async_mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +def test_mock_completion_streaming_obj_passes_prebuilt_stream_chunk_through_without_usage_chunk() -> None: + prebuilt: Final = ModelResponseStream( + model="gpt-5.4-mini", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="prebuilt"))] + ) + + assert _mock_stream_snapshots(prebuilt, 51234) == [(prebuilt.id, ("prebuilt",), None)] + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_usage_chunk() -> None: + mock_exception: Final = litellm.MockException( + status_code=500, message="boom", llm_provider="openai", model="gpt-5.4-mini" + ) + with pytest.raises(litellm.MockException): + await _async_mock_stream_snapshots(mock_exception, 51234) + + + +@contextlib.contextmanager +def _recording_hidden_params_at_submit(submit_target: str) -> "Iterator[queue.SimpleQueue[dict[str, object]]]": + seen: Final = queue.SimpleQueue() + + def record_submit(_fn, *args, **_kwargs): + response: Final = next(arg for arg in args if isinstance(arg, litellm.ModelResponse)) + seen.put(dict(response._hidden_params)) + return MagicMock() + + with patch(submit_target, side_effect=record_submit): + yield seen + + +@pytest.mark.asyncio +async def test_acompletion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(monkeypatch): + monkeypatch.setattr(litellm, "success_callback", [lambda kwargs, response, start_time, end_time: None]) + with _recording_hidden_params_at_submit("litellm.litellm_core_utils.litellm_logging.executor.submit") as seen: + await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + num_retries=0, + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"] + + +def test_completion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(): + with _recording_hidden_params_at_submit("litellm.utils.executor.submit") as seen: + litellm.completion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx index 1eeebe4ebba..118a8655d5a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx @@ -54,7 +54,7 @@ function ResourceBadge({ fallback, }: { resource: AccessGroupResource; - href: string; + href?: string; fallback: (id: string) => string; }) { const badge = ( diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 881b0b93ff9..c303dba697b 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -473,6 +473,92 @@ it("should display 'Default Proxy Admin' for user_id when value is 'default_user }); }); +describe("entity links out of the key rows", () => { + const keyRow = async () => (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement; + + const enableColumn = async (user: ReturnType, title: string) => { + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText(title)); + await user.keyboard("{Escape}"); + }; + + const enableCreatedByColumn = (user: ReturnType) => enableColumn(user, "Created By"); + + it("points the User and Team cells at their detail pages", async () => { + renderWithProviders(); + + const row = await keyRow(); + expect(within(row).getByRole("link", { name: "user@example.com" })).toHaveAttribute( + "href", + "/ui/users?user=user-1", + ); + expect(within(row).getByRole("link", { name: "Test Team" })).toHaveAttribute("href", "/ui/teams?team=team-1"); + }); + + it("points the Organization cell at the org's detail page", async () => { + mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, org_id: "org-1" }])); + const user = userEvent.setup(); + renderWithProviders(); + await enableColumn(user, "Organization"); + + const row = await keyRow(); + expect(within(row).getByRole("link", { name: "Test Organization" })).toHaveAttribute( + "href", + "/ui/organizations?org=org-1", + ); + }); + + it("points the Created By cell at the creator's detail page", async () => { + mockUseKeys.mockReturnValue( + keysResult([ + { + ...mockKey, + created_by: "creator-1", + created_by_user: { user_id: "creator-1", user_email: "creator@example.com", user_alias: "The Creator" }, + }, + ]), + ); + const user = userEvent.setup(); + renderWithProviders(); + await enableCreatedByColumn(user); + + const row = await keyRow(); + expect(within(row).getByRole("link", { name: "The Creator" })).toHaveAttribute("href", "/ui/users?user=creator-1"); + }); + + it("leaves the default_user_id placeholder unlinked even once it resolves to a named user", async () => { + const placeholder = { user_id: "default_user_id", user_email: "admin@example.com", user_alias: "Proxy Admin" }; + mockUseKeys.mockReturnValue( + keysResult([ + { + ...mockKey, + user_id: placeholder.user_id, + user_email: placeholder.user_email, + user: placeholder, + created_by: placeholder.user_id, + created_by_user: placeholder, + }, + ]), + ); + const user = userEvent.setup(); + renderWithProviders(); + await enableCreatedByColumn(user); + + const row = await keyRow(); + expect(within(row).getAllByText("Proxy Admin")).toHaveLength(2); + expect(within(row).queryByRole("link", { name: "Proxy Admin" })).not.toBeInTheDocument(); + }); + + it("leaves the litellm-dashboard session team unlinked", async () => { + mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, team_id: "litellm-dashboard" }])); + renderWithProviders(); + + const row = await keyRow(); + expect(within(row).getByText("litellm-dashboard")).toBeInTheDocument(); + expect(within(row).queryByRole("link", { name: "litellm-dashboard" })).not.toBeInTheDocument(); + }); +}); + it("should render table without crashing when models is null", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, models: null as unknown as string[] }])); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 0865fe76519..6eea77ae827 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -9,15 +9,17 @@ import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/h import { Skeleton } from "@/components/ui/skeleton"; import { DateCell, + ENTITY_CELL_TITLE_CLASSES, IdCell, IdentityCell, ModelsCell, SpendBudgetCell, StatusBadge, + UserPopoverCell, type StatusTone, } from "@/components/shared/table_cells"; +import { orgDetailHref, teamDetailHref } from "@/utils/entityLinks"; -import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; @@ -62,67 +64,6 @@ const getKeyStatus = (key: KeyResponse): KeyStatus => { }; }; -const UserPopoverCell = ({ - userAlias, - userEmail, - userId, - width, -}: { - userAlias: string | null; - userEmail: string | null; - userId: string | null; - width: number; -}) => { - const displayValue = userAlias || userEmail || userId; - const isDefaultAdmin = userId === "default_user_id"; - - const popoverContent = ( -
- {[ - { label: "User Alias", value: userAlias }, - { label: "User Email", value: userEmail }, - { label: "User ID", value: userId }, - ].map(({ label, value }) => ( -
- {label} - {value ? ( - - ) : ( - - - )} -
- ))} -
- ); - - if (isDefaultAdmin && !userAlias && !userEmail) { - return ( - - }> - - - {popoverContent} - - ); - } - - return ( - - - } - > - {displayValue || "-"} - - {popoverContent} - - ); -}; - const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => ( {label} @@ -201,12 +142,12 @@ export const getKeyTableColumns = ({ const teamId = info.getValue() as string | null; if (!teamId) return "-"; const team = allTeams.find((t) => t.team_id === teamId); - const displayValue = team?.team_alias || teamId; - const width = info.cell.column.getSize(); return ( - - {displayValue} - + ); }, }, @@ -221,12 +162,12 @@ export const getKeyTableColumns = ({ const orgId = info.getValue() as string | null; if (!orgId) return "-"; const org = organizations.find((o) => o.organization_id === orgId); - const displayValue = org?.organization_alias || orgId; - const width = info.cell.column.getSize(); return ( - - {displayValue} - + ); }, }, diff --git a/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx b/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx index 9ec24bb929b..308f722a8f2 100644 --- a/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx +++ b/ui/litellm-dashboard/src/components/common_components/DefaultProxyAdminTag.tsx @@ -1,13 +1,12 @@ import { Badge } from "@/components/ui/badge"; - -const DEFAULT_USER_ID = "default_user_id"; +import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels"; interface DefaultProxyAdminTagProps { userId: string | null | undefined; } export default function DefaultProxyAdminTag({ userId }: DefaultProxyAdminTagProps) { - if (userId === DEFAULT_USER_ID) { + if (userId === DEFAULT_PROXY_ADMIN_USER_ID) { return Default Proxy Admin; } diff --git a/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx b/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx index 9f45b05f306..99046cfa84d 100644 --- a/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx +++ b/ui/litellm-dashboard/src/components/common_components/LabeledField.tsx @@ -2,6 +2,7 @@ import React from "react"; import CopyButton from "@/components/shared/CopyButton"; import { EntityLink } from "@/components/shared/EntityLink"; import { cx } from "@/lib/cva.config"; +import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels"; import DefaultProxyAdminTag from "./DefaultProxyAdminTag"; interface LabeledFieldProps { @@ -24,7 +25,7 @@ export default function LabeledField({ defaultUserIdCheck = false, }: LabeledFieldProps) { const isEmpty = !value; - const isDefaultUser = defaultUserIdCheck && value === "default_user_id"; + const isDefaultUser = defaultUserIdCheck && value === DEFAULT_PROXY_ADMIN_USER_ID; const displayValue = isEmpty ? "-" : value; const isCopyable = copyable && !isEmpty && !isDefaultUser; const isLink = href != null && !isEmpty && !isDefaultUser; diff --git a/ui/litellm-dashboard/src/components/shared/EntityLink.test.tsx b/ui/litellm-dashboard/src/components/shared/EntityLink.test.tsx index a02f1699c30..3d6fc9a435b 100644 --- a/ui/litellm-dashboard/src/components/shared/EntityLink.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/EntityLink.test.tsx @@ -25,6 +25,12 @@ describe("EntityLink", () => { expect(push).toHaveBeenCalledWith("/ui/users?user=u1"); }); + it("renders the label as plain text when there is no href to point at", () => { + render(default_user_id); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + expect(screen.getByText("default_user_id")).toBeInTheDocument(); + }); + it("leaves modified clicks to the browser so new-tab shortcuts keep working", async () => { const user = userEvent.setup(); render(alice); diff --git a/ui/litellm-dashboard/src/components/shared/EntityLink.tsx b/ui/litellm-dashboard/src/components/shared/EntityLink.tsx index 4054b929943..4885835a42c 100644 --- a/ui/litellm-dashboard/src/components/shared/EntityLink.tsx +++ b/ui/litellm-dashboard/src/components/shared/EntityLink.tsx @@ -19,12 +19,24 @@ export function useEntityLinkClick(href: string): (e: React.MouseEvent) => void } interface EntityLinkProps { - href: string; + href?: string; className?: string; children: React.ReactNode; } export function EntityLink({ href, className, children }: EntityLinkProps) { + if (!href) { + return {children}; + } + + return ( + + {children} + + ); +} + +function LinkedEntity({ href, className, children }: EntityLinkProps & { href: string }) { const handleClick = useEntityLinkClick(href); return ( diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/UserPopoverCell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/UserPopoverCell.tsx new file mode 100644 index 00000000000..eb786e45541 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/table_cells/UserPopoverCell.tsx @@ -0,0 +1,62 @@ +"use client"; + +import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; +import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; +import { userDetailHref } from "@/utils/entityLinks"; +import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels"; + +import { IdCell } from "./id_cell"; +import { IdentityCell } from "./identity_cell"; + +export const ENTITY_CELL_TITLE_CLASSES = "font-mono text-xs font-normal"; + +interface UserPopoverCellProps { + userAlias: string | null; + userEmail: string | null; + userId: string | null; + width: number; +} + +export function UserPopoverCell({ userAlias, userEmail, userId, width }: UserPopoverCellProps) { + const displayValue = userAlias || userEmail || userId; + const isDefaultAdmin = userId === DEFAULT_PROXY_ADMIN_USER_ID; + + const popoverContent = ( +
+ {[ + { label: "User Alias", value: userAlias }, + { label: "User Email", value: userEmail }, + { label: "User ID", value: userId }, + ].map(({ label, value }) => ( +
+ {label} + {value ? ( + + ) : ( + - + )} +
+ ))} +
+ ); + + const trigger = + isDefaultAdmin && !userAlias && !userEmail ? ( + + ) : ( + + ); + + return ( + + }> + {trigger} + + {popoverContent} + + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx index 395c1815dd0..c715210fdde 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.test.tsx @@ -71,6 +71,14 @@ describe("IdCell", () => { expect(rowClick).not.toHaveBeenCalled(); }); + it("names the copy button after the field it copies", async () => { + const user = userEvent.setup(); + render(); + expect(screen.queryByRole("button", { name: "Copy ID" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Copy User Email" })); + expect(copyToClipboardMock).toHaveBeenCalledWith("alice@example.com"); + }); + it("passes dataTestId through to the id element", () => { render(); expect(screen.getByTestId("key-id-cell")).toHaveTextContent("k-1"); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx index 1b109a86106..33c7f835e64 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/id_cell.tsx @@ -15,6 +15,7 @@ interface IdCellProps { variant?: IdCellVariant; onClick?: (value: string) => void; copyable?: boolean; + copyLabel?: string; truncate?: boolean; fallback?: string; tooltip?: React.ReactNode; @@ -39,6 +40,7 @@ export function IdCell({ variant = "pill", onClick, copyable = false, + copyLabel = "Copy ID", truncate = true, fallback = "-", tooltip, @@ -80,7 +82,7 @@ export function IdCell({ {withTooltip}