mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_redis_breaker_open_silent_miss
This commit is contained in:
commit
fdd423128e
48 changed files with 1416 additions and 236 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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'."
|
||||
),
|
||||
|
|
|
|||
147
litellm/utils.py
147
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 ############################
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ function ResourceBadge({
|
|||
fallback,
|
||||
}: {
|
||||
resource: AccessGroupResource;
|
||||
href: string;
|
||||
href?: string;
|
||||
fallback: (id: string) => string;
|
||||
}) {
|
||||
const badge = (
|
||||
|
|
|
|||
|
|
@ -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<typeof userEvent.setup>, 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<typeof userEvent.setup>) => enableColumn(user, "Created By");
|
||||
|
||||
it("points the User and Team cells at their detail pages", async () => {
|
||||
renderWithProviders(<VirtualKeysTable />);
|
||||
|
||||
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(<VirtualKeysTable />);
|
||||
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(<VirtualKeysTable />);
|
||||
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(<VirtualKeysTable />);
|
||||
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(<VirtualKeysTable />);
|
||||
|
||||
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[] }]));
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
|
||||
{[
|
||||
{ label: "User Alias", value: userAlias },
|
||||
{ label: "User Email", value: userEmail },
|
||||
{ label: "User ID", value: userId },
|
||||
].map(({ label, value }) => (
|
||||
<div key={label} className="flex flex-col min-w-0">
|
||||
<span className="text-muted-foreground">{label}</span>
|
||||
{value ? (
|
||||
<IdCell value={value} variant="plain" copyable className="max-w-full" />
|
||||
) : (
|
||||
<span className="font-mono">-</span>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
);
|
||||
|
||||
if (isDefaultAdmin && !userAlias && !userEmail) {
|
||||
return (
|
||||
<HoverCard>
|
||||
<HoverCardTrigger render={<span className="cursor-default" />}>
|
||||
<DefaultProxyAdminTag userId={userId} />
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<HoverCard>
|
||||
<HoverCardTrigger
|
||||
render={
|
||||
<span
|
||||
className="font-mono text-xs truncate block cursor-default"
|
||||
style={{ maxWidth: width, overflow: "hidden" }}
|
||||
/>
|
||||
}
|
||||
>
|
||||
{displayValue || "-"}
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
};
|
||||
|
||||
const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => (
|
||||
<span className="flex items-center gap-1">
|
||||
{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 (
|
||||
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
|
||||
{displayValue}
|
||||
</span>
|
||||
<IdentityCell
|
||||
title={team?.team_alias || teamId}
|
||||
titleClassName={ENTITY_CELL_TITLE_CLASSES}
|
||||
href={teamDetailHref(teamId)}
|
||||
/>
|
||||
);
|
||||
},
|
||||
},
|
||||
|
|
@ -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 (
|
||||
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
|
||||
{displayValue}
|
||||
</span>
|
||||
<IdentityCell
|
||||
title={org?.organization_alias || orgId}
|
||||
titleClassName={ENTITY_CELL_TITLE_CLASSES}
|
||||
href={orgDetailHref(orgId)}
|
||||
/>
|
||||
);
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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 <Badge variant="secondary">Default Proxy Admin</Badge>;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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(<EntityLink>default_user_id</EntityLink>);
|
||||
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(<EntityLink href="/ui/users?user=u1">alice</EntityLink>);
|
||||
|
|
|
|||
|
|
@ -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 <span className={cn("inline-block min-w-0 max-w-full truncate font-semibold", className)}>{children}</span>;
|
||||
}
|
||||
|
||||
return (
|
||||
<LinkedEntity href={href} className={className}>
|
||||
{children}
|
||||
</LinkedEntity>
|
||||
);
|
||||
}
|
||||
|
||||
function LinkedEntity({ href, className, children }: EntityLinkProps & { href: string }) {
|
||||
const handleClick = useEntityLinkClick(href);
|
||||
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
|
||||
{[
|
||||
{ label: "User Alias", value: userAlias },
|
||||
{ label: "User Email", value: userEmail },
|
||||
{ label: "User ID", value: userId },
|
||||
].map(({ label, value }) => (
|
||||
<div key={label} className="flex flex-col min-w-0">
|
||||
<span className="text-muted-foreground">{label}</span>
|
||||
{value ? (
|
||||
<IdCell value={value} variant="plain" copyable copyLabel={`Copy ${label}`} className="max-w-full" />
|
||||
) : (
|
||||
<span className="font-mono">-</span>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
);
|
||||
|
||||
const trigger =
|
||||
isDefaultAdmin && !userAlias && !userEmail ? (
|
||||
<DefaultProxyAdminTag userId={userId} />
|
||||
) : (
|
||||
<IdentityCell
|
||||
title={displayValue || "-"}
|
||||
titleClassName={ENTITY_CELL_TITLE_CLASSES}
|
||||
href={userId ? userDetailHref(userId) : undefined}
|
||||
/>
|
||||
);
|
||||
|
||||
return (
|
||||
<HoverCard>
|
||||
<HoverCardTrigger render={<span className="block" style={{ maxWidth: width, overflow: "hidden" }} />}>
|
||||
{trigger}
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
}
|
||||
|
|
@ -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(<IdCell value="alice@example.com" copyable copyLabel="Copy User Email" />);
|
||||
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(<IdCell value="k-1" dataTestId="key-id-cell" />);
|
||||
expect(screen.getByTestId("key-id-cell")).toHaveTextContent("k-1");
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
<button
|
||||
type="button"
|
||||
aria-label="Copy ID"
|
||||
aria-label={copyLabel}
|
||||
className="shrink-0 cursor-pointer text-muted-foreground hover:text-foreground"
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
|
|
|
|||
|
|
@ -13,3 +13,4 @@ export { ModelsCell } from "./models_cell";
|
|||
export { MoneyCell } from "./money_cell";
|
||||
export { SpendBudgetCell } from "./spend_budget_cell";
|
||||
export { StatusBadge, type StatusTone } from "./status_badge";
|
||||
export { UserPopoverCell, ENTITY_CELL_TITLE_CLASSES } from "./UserPopoverCell";
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
|
|||
useKeys: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) }));
|
||||
|
||||
vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({
|
||||
getModelDisplayName: vi.fn((model: string) => model),
|
||||
}));
|
||||
|
|
@ -384,4 +386,66 @@ describe("TeamVirtualKeysTable", () => {
|
|||
expect(screen.getByText("Key Info View")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe("entity links out of the key rows", () => {
|
||||
const renderRow = async (key: KeyResponse, organization: Organization | null = null) => {
|
||||
mockUseKeys.mockReturnValue({
|
||||
data: { keys: [key], total_count: 1, current_page: 1, total_pages: 1 } as KeysResponse,
|
||||
isPending: false,
|
||||
isFetching: false,
|
||||
refetch: vi.fn(),
|
||||
} as any);
|
||||
renderWithProviders(<TeamVirtualKeysTable {...defaultProps} organization={organization} />);
|
||||
await screen.findByText(key.key_alias as string);
|
||||
return screen.getByRole("row", { name: new RegExp(key.key_alias as string) });
|
||||
};
|
||||
|
||||
it("points the Organization ID cell at the org's detail page", async () => {
|
||||
const row = await renderRow(createMockKey({ organization_id: null }), mockOrganization);
|
||||
expect(within(row).getByRole("link", { name: "org-123" })).toHaveAttribute(
|
||||
"href",
|
||||
"/ui/organizations?org=org-123",
|
||||
);
|
||||
});
|
||||
|
||||
it("points the User Email and User ID cells at the owning user's detail page", async () => {
|
||||
const row = await renderRow(
|
||||
createMockKey({ user_id: "user-1", user: { user_id: "user-1", user_email: "alice@example.com" } }),
|
||||
);
|
||||
expect(within(row).getByRole("link", { name: "alice@example.com" })).toHaveAttribute(
|
||||
"href",
|
||||
"/ui/users?user=user-1",
|
||||
);
|
||||
expect(within(row).getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1");
|
||||
});
|
||||
|
||||
it("points the Created By cell at the creator's detail page", async () => {
|
||||
const row = await renderRow(
|
||||
createMockKey({
|
||||
created_by: "creator-1",
|
||||
created_by_user: { user_id: "creator-1", user_email: "creator@example.com", user_alias: "The Creator" },
|
||||
}),
|
||||
);
|
||||
expect(within(row).getByRole("link", { name: "The Creator" })).toHaveAttribute(
|
||||
"href",
|
||||
"/ui/users?user=creator-1",
|
||||
);
|
||||
});
|
||||
|
||||
it("leaves the default_user_id placeholder unlinked in the User ID and Created By cells", async () => {
|
||||
const placeholder = { user_id: "default_user_id", user_email: "admin@example.com", user_alias: "Proxy Admin" };
|
||||
const ownedAndCreatedByPlaceholder = {
|
||||
user_id: placeholder.user_id,
|
||||
user: placeholder,
|
||||
created_by: placeholder.user_id,
|
||||
created_by_user: placeholder,
|
||||
};
|
||||
const row = await renderRow(createMockKey(ownedAndCreatedByPlaceholder));
|
||||
expect(within(row).getByText("Default Proxy Admin")).toBeInTheDocument();
|
||||
expect(within(row).getByText("Proxy Admin")).toBeInTheDocument();
|
||||
expect(within(row).queryByRole("link", { name: "Proxy Admin" })).not.toBeInTheDocument();
|
||||
expect(within(row).queryByRole("link", { name: placeholder.user_email })).not.toBeInTheDocument();
|
||||
expect(within(row).queryByRole("link", { name: "Default Proxy Admin" })).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,8 +1,14 @@
|
|||
"use client";
|
||||
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import CopyButton from "@/components/shared/CopyButton";
|
||||
import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells";
|
||||
import {
|
||||
DateCell,
|
||||
ENTITY_CELL_TITLE_CLASSES,
|
||||
IdCell,
|
||||
IdentityCell,
|
||||
MoneyCell,
|
||||
UserPopoverCell,
|
||||
} from "@/components/shared/table_cells";
|
||||
import {
|
||||
DataTable,
|
||||
DataTableFilterDrawer,
|
||||
|
|
@ -11,8 +17,9 @@ import {
|
|||
DataTableToolbar,
|
||||
} from "@/components/shared/DataTable";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { orgDetailHref, userDetailHref } from "@/utils/entityLinks";
|
||||
import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels";
|
||||
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
|
||||
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
|
||||
import { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
|
||||
|
|
@ -168,7 +175,15 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
header: "Organization ID",
|
||||
size: 140,
|
||||
enableSorting: false,
|
||||
cell: (info) => (info.getValue() ? info.renderValue() : "-"),
|
||||
cell: (info) => {
|
||||
const orgId = info.getValue() as string | null;
|
||||
if (!orgId) return "-";
|
||||
return (
|
||||
<SimpleTooltip content={orgId}>
|
||||
<IdentityCell title={orgId} titleClassName={ENTITY_CELL_TITLE_CLASSES} href={orgDetailHref(orgId)} />
|
||||
</SimpleTooltip>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "user_email",
|
||||
|
|
@ -179,9 +194,14 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
cell: (info) => {
|
||||
const user = info.getValue() as { user_email?: string } | undefined;
|
||||
const value = user?.user_email;
|
||||
const userId = info.row.original.user_id;
|
||||
return (
|
||||
<SimpleTooltip content={value}>
|
||||
<span className="block max-w-full truncate font-mono text-xs">{value ?? "-"}</span>
|
||||
<IdentityCell
|
||||
title={value ?? "-"}
|
||||
titleClassName={ENTITY_CELL_TITLE_CLASSES}
|
||||
href={value && userId ? userDetailHref(userId) : undefined}
|
||||
/>
|
||||
</SimpleTooltip>
|
||||
);
|
||||
},
|
||||
|
|
@ -194,10 +214,16 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
enableSorting: false,
|
||||
cell: (info) => {
|
||||
const userId = info.getValue() as string | null;
|
||||
const displayValue = userId === "default_user_id" ? "Default Proxy Admin" : userId;
|
||||
if (userId === DEFAULT_PROXY_ADMIN_USER_ID) {
|
||||
return <DefaultProxyAdminTag userId={userId} />;
|
||||
}
|
||||
return (
|
||||
<SimpleTooltip content={displayValue}>
|
||||
<span className="block max-w-full truncate font-mono text-xs">{displayValue ?? "-"}</span>
|
||||
<SimpleTooltip content={userId}>
|
||||
<IdentityCell
|
||||
title={userId ?? "-"}
|
||||
titleClassName={ENTITY_CELL_TITLE_CLASSES}
|
||||
href={userId ? userDetailHref(userId) : undefined}
|
||||
/>
|
||||
</SimpleTooltip>
|
||||
);
|
||||
},
|
||||
|
|
@ -221,53 +247,13 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
const userId = info.getValue() as string | null;
|
||||
if (!userId) return "-";
|
||||
const { created_by_user } = info.row.original;
|
||||
const userAlias = created_by_user?.user_alias ?? null;
|
||||
const userEmail = created_by_user?.user_email ?? null;
|
||||
const isDefaultAdmin = userId === "default_user_id";
|
||||
const displayValue = userAlias || userEmail || userId;
|
||||
|
||||
const popoverContent = (
|
||||
<div className="flex min-w-[200px] max-w-[300px] flex-col gap-2 text-xs">
|
||||
{[
|
||||
{ label: "User Alias", value: userAlias },
|
||||
{ label: "User Email", value: userEmail },
|
||||
{ label: "User ID", value: userId },
|
||||
].map(({ label, value }) => (
|
||||
<div key={label} className="flex flex-col min-w-0">
|
||||
<span className="text-muted-foreground">{label}</span>
|
||||
{value ? (
|
||||
<span className="flex items-center gap-1">
|
||||
<span className="min-w-0 flex-1 truncate font-mono text-xs">{value}</span>
|
||||
<CopyButton value={value} label={`Copy ${label}`} />
|
||||
</span>
|
||||
) : (
|
||||
<span className="font-mono">-</span>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
);
|
||||
|
||||
if (isDefaultAdmin && !userAlias && !userEmail) {
|
||||
return (
|
||||
<HoverCard>
|
||||
<HoverCardTrigger render={<span className="cursor-default" />}>
|
||||
<DefaultProxyAdminTag userId={userId} />
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<HoverCard>
|
||||
<HoverCardTrigger
|
||||
render={<span className="block max-w-full cursor-default truncate font-mono text-xs" />}
|
||||
>
|
||||
{displayValue}
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
|
||||
</HoverCard>
|
||||
<UserPopoverCell
|
||||
userAlias={created_by_user?.user_alias ?? null}
|
||||
userEmail={created_by_user?.user_email ?? null}
|
||||
userId={userId}
|
||||
width={130}
|
||||
/>
|
||||
);
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -2,7 +2,29 @@ import { describe, expect, it, vi } from "vitest";
|
|||
|
||||
vi.mock("@/components/networking", () => ({ serverRootPath: "" }));
|
||||
|
||||
import { modelGroupHref } from "./entityLinks";
|
||||
import { modelGroupHref, teamDetailHref, userDetailHref } from "./entityLinks";
|
||||
|
||||
describe("userDetailHref", () => {
|
||||
it("targets the users page filtered to the encoded user id", () => {
|
||||
expect(userDetailHref("user-1")).toMatch(/\/users\?user=user-1$/);
|
||||
expect(userDetailHref("a b/c")).toMatch(/\?user=a%20b%2Fc$/);
|
||||
});
|
||||
|
||||
it("returns no href for the proxy admin placeholder, which has no user page", () => {
|
||||
expect(userDetailHref("default_user_id")).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("teamDetailHref", () => {
|
||||
it("targets the teams page filtered to the encoded team id", () => {
|
||||
expect(teamDetailHref("team-1")).toMatch(/\/teams\?team=team-1$/);
|
||||
expect(teamDetailHref("a b/c")).toMatch(/\?team=a%20b%2Fc$/);
|
||||
});
|
||||
|
||||
it("returns no href for the Admin UI session team, which has no team page", () => {
|
||||
expect(teamDetailHref("litellm-dashboard")).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("modelGroupHref", () => {
|
||||
it("targets the models page filtered to the encoded model group", () => {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { DEFAULT_PROXY_ADMIN_USER_ID, UI_TEAM_ID } from "@/utils/sentinels";
|
||||
import { uiHref } from "@/utils/uiHref";
|
||||
|
||||
const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
|
||||
|
|
@ -6,7 +7,8 @@ const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
|
|||
"no-default-models",
|
||||
]);
|
||||
|
||||
export function teamDetailHref(teamId: string): string {
|
||||
export function teamDetailHref(teamId: string): string | undefined {
|
||||
if (teamId === UI_TEAM_ID) return undefined;
|
||||
return `${uiHref("teams")}?team=${encodeURIComponent(teamId)}`;
|
||||
}
|
||||
|
||||
|
|
@ -14,7 +16,8 @@ export function keyDetailHref(keyToken: string): string {
|
|||
return `${uiHref("api-keys")}?key=${encodeURIComponent(keyToken)}`;
|
||||
}
|
||||
|
||||
export function userDetailHref(userId: string): string {
|
||||
export function userDetailHref(userId: string): string | undefined {
|
||||
if (userId === DEFAULT_PROXY_ADMIN_USER_ID) return undefined;
|
||||
return `${uiHref("users")}?user=${encodeURIComponent(userId)}`;
|
||||
}
|
||||
|
||||
|
|
|
|||
3
ui/litellm-dashboard/src/utils/sentinels.ts
Normal file
3
ui/litellm-dashboard/src/utils/sentinels.ts
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
export const DEFAULT_PROXY_ADMIN_USER_ID = "default_user_id";
|
||||
|
||||
export const UI_TEAM_ID = "litellm-dashboard";
|
||||
4
uv.lock
generated
4
uv.lock
generated
|
|
@ -10,7 +10,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-09-06T00:40:30.433549Z"
|
||||
exclude-newer = "2026-09-07T23:09:03.362777Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -4772,7 +4772,7 @@ source = { editable = "enterprise" }
|
|||
|
||||
[[package]]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.95"
|
||||
version = "0.4.96"
|
||||
source = { editable = "litellm-proxy-extras" }
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue