diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index bb04563c1a8..206bb809e0c 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -80,7 +80,7 @@ jobs: LITELLM_IMAGE: litellm-image-scan:${{ github.sha }} run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v # Scans the whole shipped artifact: OS/apk plus every language package # baked into the image, including ones no lockfile declares (e.g. prisma's @@ -124,7 +124,7 @@ jobs: LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }} run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v migrations-image: name: migrations-image @@ -185,7 +185,7 @@ jobs: LITELLM_COMPONENT_PORT: "4000" run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v ui-image: name: ui-image diff --git a/Dockerfile b/Dockerfile index b3ee85e9ed1..0a92aa9a68c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ @@ -101,6 +103,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root +# The base image only configures Chainguard's authenticated apk repo, which +# requires an enterprise subscription. Add the public Wolfi repo so `apk add` +# also works for anyone installing extra packages into a running container. +# https://github.com/BerriAI/litellm/issues/33518 +RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories + # node (without npm) is required by the prisma CLI at runtime RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile diff --git a/README.md b/README.md index 68aaa09ec98..92757fcbbc1 100644 --- a/README.md +++ b/README.md @@ -354,6 +354,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ | [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | | | [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | | | [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | | +| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | +| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ | | [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | | | [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | | | [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/backend/Dockerfile b/backend/Dockerfile index aa01b9fba8b..622fedcd70d 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3.13 # Stage 2 — copy source and install the project + workspace members. @@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 551423c7707..da788bf1ce3 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 16171 + "limit": 14076 }, "reportArgumentType": { - "limit": 2224 + "limit": 2216 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5199 + "limit": 4128 }, "reportFunctionMemberAccess": { "limit": 7 @@ -42,7 +42,7 @@ "limit": 12 }, "reportIndexIssue": { - "limit": 35 + "limit": 25 }, "reportInvalidTypeForm": { "limit": 34 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5611 + "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15348 + "limit": 15306 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,25 +105,25 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38465 + "limit": 38350 }, "reportUnknownParameterType": { - "limit": 19663 + "limit": 19626 }, "reportUnknownVariableType": { - "limit": 30064 + "limit": 29890 }, "reportUnnecessaryCast": { "limit": 111 }, "reportUnnecessaryComparison": { - "limit": 695 + "limit": 692 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 828 + "limit": 826 }, "reportUntypedBaseClass": { "limit": 0 @@ -141,6 +141,6 @@ "limit": 543 }, "reportUnusedVariable": { - "limit": 139 + "limit": 137 } } diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index c1348f68231..e9ad2849bb2 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2221435a83a..edf20e8bbff 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13; \ fi diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 3b09dc9272e..354a6ed2fd0 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -5,7 +5,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast +from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -87,7 +87,7 @@ class CheckBatchCost: return self.batch_processed_support_confirmed = True - async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: + async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]: """ Look up user email and key alias by user_id for enriching the S3 callback metadata. Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). @@ -97,8 +97,10 @@ class CheckBatchCost: if not user_id: return {} try: - user_row = await self.prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} + user_row: prisma_models.LiteLLM_UserTable | None = ( + await self.prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) ) if user_row is None: return {} @@ -115,8 +117,10 @@ class CheckBatchCost: if not api_key: return None try: - key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": api_key} + key_row: prisma_models.LiteLLM_VerificationToken | None = ( + await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) ) return getattr(key_row, "key_alias", None) if key_row is not None else None except Exception as e: @@ -128,8 +132,10 @@ class CheckBatchCost: if not team_id: return None try: - team_row = await self.prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row: prisma_models.LiteLLM_TeamTable | None = ( + await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) return getattr(team_row, "team_alias", None) if team_row is not None else None except Exception as e: @@ -138,7 +144,7 @@ class CheckBatchCost: async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str - ) -> Dict[str, Any]: + ) -> dict[str, object]: """ Rebuild the spend-tracking metadata for the key, team, and tags that created the batch so the batch-cost spend log is attributed the same way a non-batch request @@ -152,7 +158,7 @@ class CheckBatchCost: team_id = getattr(job, "team_id", None) request_tags = getattr(job, "request_tags", None) - metadata: Dict[str, Any] = { + metadata: dict[str, object] = { "user_api_key_user_id": job.created_by, "user_api_key": api_key, "user_api_key_team_id": team_id, diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 570b306d6df..5cfcf6129f0 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -182,6 +182,10 @@ class _ManagedObjectTableActions(Protocol): async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... +class _SchedulerWithJobLookup(Protocol): + def get_job(self, job_id: str) -> object: ... + + class _CursorPageArgs(TypedDict, total=False): cursor: Mapping[str, str] skip: int @@ -853,7 +857,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]: """ Gets file ids from responses API input. @@ -878,7 +882,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check for direct input_file type if item.get("type") == "input_file": file_id = item.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: file_ids.append(file_id) # Check for input_file in content array @@ -887,7 +891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for content_item in content: if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") - if file_id: + if isinstance(file_id, str) and file_id: file_ids.append(file_id) return file_ids @@ -1227,7 +1231,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: - file_id_value = getattr(response, file_attr, None) + file_id_value: str | None = getattr(response, file_attr, None) if file_id_value and model_id: decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: @@ -1496,7 +1500,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): import litellm.proxy.proxy_server as proxy_server_module # Check if the scheduler has the batch cost checking job registered - scheduler = getattr(proxy_server_module, "scheduler", None) + scheduler: Final[_SchedulerWithJobLookup | None] = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False @@ -1542,7 +1546,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) MAX_MATCHES_TO_RETURN = 10 - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where={ "file_purpose": "batch", "batch_processed": False, @@ -1552,11 +1556,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): order={"created_at": "desc"}, ) - referencing_batches = [] + referencing_batches: Final[list[dict[str, object]]] = [] for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object + decoded_file_object = _decode_json_blob(batch.file_object) + batch_data: Mapping[str, object] = ( + decoded_file_object if isinstance(decoded_file_object, Mapping) else {} + ) # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index cac98b69793..8360c0a077d 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.62" +version = "0.1.63" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.62" +version = "0.1.63" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/helm/litellm-helm/Chart.yaml b/helm/litellm-helm/Chart.yaml index 3959d85edf3..a3cb388ffc6 100644 --- a/helm/litellm-helm/Chart.yaml +++ b/helm/litellm-helm/Chart.yaml @@ -18,7 +18,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.2 +version: 1.1.3 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index b242373de5d..bf4089404db 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented): | `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` | | `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A | | `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A | -| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A | +| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A | | `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | | `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | | `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | @@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://-litellm:4000` The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey` was not provided to the helm command line, the `masterkey` is a randomly generated string in the `sk-...` format stored in the `-litellm-masterkey` Kubernetes Secret. +The key is generated once on the first install; later `helm upgrade` runs reuse the +value already in that Secret, so upgrading never rotates the master key. ```bash kubectl -n litellm get secret -litellm-masterkey -o jsonpath="{.data.masterkey}" diff --git a/helm/litellm-helm/templates/secret-masterkey.yaml b/helm/litellm-helm/templates/secret-masterkey.yaml index 7c8560cc2cc..60ab4e74c6b 100644 --- a/helm/litellm-helm/templates/secret-masterkey.yaml +++ b/helm/litellm-helm/templates/secret-masterkey.yaml @@ -1,9 +1,11 @@ {{- if not .Values.masterkeySecretName }} -{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }} +{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }} +{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }} +{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }} apiVersion: v1 kind: Secret metadata: - name: {{ include "litellm.fullname" . }}-masterkey + name: {{ $secretName }} data: masterkey: {{ $masterkey | b64enc }} type: Opaque diff --git a/helm/litellm-helm/tests/hpa_tests.yaml b/helm/litellm-helm/tests/hpa_tests.yaml index ec18c3591d3..cd062dd5971 100644 --- a/helm/litellm-helm/tests/hpa_tests.yaml +++ b/helm/litellm-helm/tests/hpa_tests.yaml @@ -1,4 +1,4 @@ -suite: "hpa with behavior" +suite: "hpa" templates: - hpa.yaml tests: @@ -23,14 +23,44 @@ tests: - equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 } - equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 } ---- -suite: "hpa without behavior" -templates: - - hpa.yaml -tests: - it: "does not render behavior when not set" set: autoscaling.enabled: true asserts: - isKind: { of: HorizontalPodAutoscaler } - isNull: { path: spec.behavior } + + - it: "scales on cpu at the documented 60 percent by default" + set: + autoscaling.enabled: true + asserts: + - isKind: { of: HorizontalPodAutoscaler } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + - equal: { path: "spec.metrics[0].resource.target.type", value: Utilization } + - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 } + + - it: "does not scale on memory by default" + set: + autoscaling.enabled: true + asserts: + - lengthEqual: { path: spec.metrics, count: 1 } + + - it: "honours an explicit cpu target override" + set: + autoscaling.enabled: true + autoscaling.targetCPUUtilizationPercentage: 75 + asserts: + - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 } + + - it: "renders a memory metric only when a memory target is set" + set: + autoscaling.enabled: true + autoscaling.targetMemoryUtilizationPercentage: 80 + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - equal: { path: "spec.metrics[1].resource.name", value: memory } + - equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 } + + - it: "renders no hpa when autoscaling is disabled" + asserts: + - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/tests/masterkey-secret_tests.yaml b/helm/litellm-helm/tests/masterkey-secret_tests.yaml index bbbade9d802..296f26755b8 100644 --- a/helm/litellm-helm/tests/masterkey-secret_tests.yaml +++ b/helm/litellm-helm/tests/masterkey-secret_tests.yaml @@ -15,6 +15,53 @@ tests: # Note: The masterkey is generated as "sk-<18-random-chars>" in plain text, # but stored as base64 encoded in Kubernetes secret (requirement). # "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern. + - it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade + template: secret-masterkey.yaml + set: + masterkeySecretName: "" + kubernetesProvider: + scheme: + "v1/Secret": + gvr: + version: "v1" + resource: "secrets" + namespaced: true + objects: + - kind: Secret + apiVersion: v1 + metadata: + name: RELEASE-NAME-litellm-masterkey + namespace: NAMESPACE + data: + masterkey: c2stZXhpc3Rpbmcta2V5 + asserts: + - equal: + path: data.masterkey + value: c2stZXhpc3Rpbmcta2V5 + - it: should let an explicit masterkey value override the one already stored in the cluster + template: secret-masterkey.yaml + set: + masterkeySecretName: "" + masterkey: sk-explicit + kubernetesProvider: + scheme: + "v1/Secret": + gvr: + version: "v1" + resource: "secrets" + namespaced: true + objects: + - kind: Secret + apiVersion: v1 + metadata: + name: RELEASE-NAME-litellm-masterkey + namespace: NAMESPACE + data: + masterkey: c2stZXhpc3Rpbmcta2V5 + asserts: + - equal: + path: data.masterkey + value: c2stZXhwbGljaXQ= - it: should not create a secret if masterkeySecretName is set template: secret-masterkey.yaml set: diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index f8df98de102..637be2322e3 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -200,7 +200,16 @@ autoscaling: enabled: false minReplicas: 1 maxReplicas: 100 - targetCPUUtilizationPercentage: 80 + # 60 is the documented recommendation. See "Recommended Machine Specifications" + # in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe + # above only after up to failureThreshold x periodSeconds = 300 seconds, so a target + # high enough to trip near saturation adds capacity minutes after it was needed. + targetCPUUtilizationPercentage: 60 + # Deliberately left unset rather than given a value. The prisma query engine's + # resident memory is a high-water mark that ratchets to the pod's worst-ever write + # and is never returned, so a memory target reads the largest write a pod ever did + # rather than what it is doing now, and replicas ratchet up without scaling back in. + # Memory is a floor to provision under 'resources', not a signal to scale on. # targetMemoryUtilizationPercentage: 80 # behavior: {} diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 5c0431fc0bd..0db2f0b3d43 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: backend spec: + {{- with .Values.backend.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.backend.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index d5363d0096e..5030ba2c9dc 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: gateway spec: + {{- with .Values.gateway.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.gateway.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index 9cd8397f794..8d33081e72f 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -7,6 +7,8 @@ # # Running this pre-upgrade closes the window where new application pods would # otherwise serve traffic against the previous release's unmigrated schema. +# Argo CD users can swap the Helm hook for a PreSync hook through +# `migrationJob.hooks`, which re-runs the Job on every sync. apiVersion: batch/v1 kind: Job metadata: @@ -14,10 +16,18 @@ metadata: labels: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: migrations + {{- if or .Values.migrationJob.hooks.helm.enabled .Values.migrationJob.hooks.argocd.enabled }} annotations: + {{- if .Values.migrationJob.hooks.helm.enabled }} helm.sh/hook: pre-install,pre-upgrade helm.sh/hook-delete-policy: before-hook-creation - helm.sh/hook-weight: "0" + helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "0" | quote }} + {{- end }} + {{- if .Values.migrationJob.hooks.argocd.enabled }} + argocd.argoproj.io/hook: PreSync + argocd.argoproj.io/hook-delete-policy: BeforeHookCreation + {{- end }} + {{- end }} spec: backoffLimit: {{ .Values.migrationJob.backoffLimit }} ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index 91d6de39ea6..b992b347bad 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -7,6 +7,10 @@ metadata: {{- include "litellm.commonLabels" . | nindent 4 }} app.kubernetes.io/component: ui spec: + {{- with .Values.ui.strategy }} + strategy: + {{- toYaml . | nindent 4 }} + {{- end }} selector: matchLabels: {{- include "litellm.ui.selectorLabels" . | nindent 6 }} diff --git a/helm/litellm/tests/migration_job_hooks_tests.yaml b/helm/litellm/tests/migration_job_hooks_tests.yaml new file mode 100644 index 00000000000..650d2700429 --- /dev/null +++ b/helm/litellm/tests/migration_job_hooks_tests.yaml @@ -0,0 +1,63 @@ +suite: test migrations Job hook annotations +templates: + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: runs as a Helm pre-install / pre-upgrade hook by default + asserts: + - equal: + path: metadata.annotations["helm.sh/hook"] + value: pre-install,pre-upgrade + - equal: + path: metadata.annotations["helm.sh/hook-delete-policy"] + value: before-hook-creation + - equal: + path: metadata.annotations["helm.sh/hook-weight"] + value: "0" + - notExists: + path: metadata.annotations["argocd.argoproj.io/hook"] + + - it: adds the Argo CD PreSync hook when asked + set: + migrationJob.hooks.argocd.enabled: true + asserts: + - equal: + path: metadata.annotations["argocd.argoproj.io/hook"] + value: PreSync + - equal: + path: metadata.annotations["argocd.argoproj.io/hook-delete-policy"] + value: BeforeHookCreation + + - it: drops the Helm hook so Argo CD owns the Job + set: + migrationJob.hooks.argocd.enabled: true + migrationJob.hooks.helm.enabled: false + asserts: + - equal: + path: metadata.annotations["argocd.argoproj.io/hook"] + value: PreSync + - notExists: + path: metadata.annotations["helm.sh/hook"] + - notExists: + path: metadata.annotations["helm.sh/hook-delete-policy"] + - notExists: + path: metadata.annotations["helm.sh/hook-weight"] + + - it: renders an ordinary Job when both hooks are disabled + set: + migrationJob.hooks.helm.enabled: false + asserts: + - notExists: + path: metadata.annotations + - equal: + path: kind + value: Job + + - it: honours a custom Helm hook weight + set: + migrationJob.hooks.helm.weight: "-5" + asserts: + - equal: + path: metadata.annotations["helm.sh/hook-weight"] + value: "-5" diff --git a/helm/litellm/tests/rollout_strategy_tests.yaml b/helm/litellm/tests/rollout_strategy_tests.yaml new file mode 100644 index 00000000000..b12e2073c7c --- /dev/null +++ b/helm/litellm/tests/rollout_strategy_tests.yaml @@ -0,0 +1,66 @@ +suite: test rolling update strategy on the component deployments +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: leaves the strategy to Kubernetes defaults when unset + asserts: + - notExists: + path: spec.strategy + + - it: renders the configured strategy on each deployment + set: + gateway.strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + backend.strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: "25%" + maxSurge: 2 + ui.strategy: + type: Recreate + asserts: + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 1 + template: gateway/deployment.yaml + - equal: + path: spec.strategy + value: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 25% + maxSurge: 2 + template: backend/deployment.yaml + - equal: + path: spec.strategy + value: + type: Recreate + template: ui/deployment.yaml + + - it: keeps a component on the cluster default when only another one sets a strategy + set: + gateway.strategy: + type: Recreate + asserts: + - equal: + path: spec.strategy.type + value: Recreate + template: gateway/deployment.yaml + - notExists: + path: spec.strategy + template: backend/deployment.yaml + - notExists: + path: spec.strategy + template: ui/deployment.yaml diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index d0c80fd6f6f..378c3b7a618 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -75,6 +75,22 @@ serviceAccounts: # generate` — the migration engine doesn't need the generated client. migrationJob: enabled: true + # Which controller is responsible for running the Job. + # + # `helm.enabled` renders the Helm pre-install / pre-upgrade hook, so the Job + # runs whenever `helm upgrade` sees a change to apply. `argocd.enabled` + # renders an Argo CD PreSync hook instead, which runs the Job on every sync + # even when the rendered manifests are unchanged: the way to re-run + # migrations on demand from a GitOps pipeline. Turning the Helm hook off + # while the Argo CD hook is on leaves the Job out of Helm's own upgrade + # path, which is what Argo CD users want since Argo, not Helm, applies the + # manifests. + hooks: + helm: + enabled: true + weight: "0" + argocd: + enabled: false backoffLimit: 4 ttlSecondsAfterFinished: 120 # Wall-clock budget for the whole Job, shared across every `backoffLimit` @@ -257,6 +273,15 @@ gateway: initialDelaySeconds: 5 periodSeconds: 10 timeoutSeconds: 10 + # Rolling update tuning for the gateway Deployment. Empty by default, so + # Kubernetes applies its own RollingUpdate defaults (25% maxSurge / + # 25% maxUnavailable). Example, for a surge-only rollout behind a load + # balancer that must never lose capacity: + # type: RollingUpdate + # rollingUpdate: + # maxUnavailable: 0 + # maxSurge: 1 + strategy: {} # Optional startupProbe. Empty by default, so existing installs are unchanged # and liveness/readiness apply from container start. Set it to gate # liveness/readiness until a slow cold start finishes — a high failureThreshold @@ -369,6 +394,8 @@ backend: initialDelaySeconds: 5 periodSeconds: 10 timeoutSeconds: 10 + # Same shape as gateway.strategy. + strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} hpa: @@ -433,6 +460,8 @@ ui: httpGet: { path: /, port: http } initialDelaySeconds: 2 periodSeconds: 10 + # Same shape as gateway.strategy. + strategy: {} # Optional startupProbe; same shape as gateway.startupProbe. Empty by default. startupProbe: {} hpa: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql new file mode 100644 index 00000000000..90b21205310 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql @@ -0,0 +1,3 @@ +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "router_names" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[]; + +ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "router_name" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a9c468c3e9d..b44935ce3dd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1533,7 +1533,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1557,6 +1558,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index d5741d479bf..0944f99ad54 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.91" +version = "0.4.92" 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.91" +version = "0.4.92" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 1447e05fdf7..4eeececdb7e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -659,6 +659,8 @@ aiml_models: Set = set() deepgram_models: Set = set() elevenlabs_models: Set = set() dashscope_models: Set = set() +qwencloud_models: Set = set() +qwen_ai_platform_models: Set = set() moonshot_models: Set = set() publicai_models: Set = set() darkbloom_models: Set = set() @@ -909,6 +911,10 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None: heroku_models.add(key) elif value.get("litellm_provider") == "dashscope": dashscope_models.add(key) + elif value.get("litellm_provider") == "qwencloud": + qwencloud_models.add(key) + elif value.get("litellm_provider") == "qwen_ai_platform": + qwen_ai_platform_models.add(key) elif value.get("litellm_provider") == "modelscope": modelscope_models.add(key) elif value.get("litellm_provider") == "moonshot": @@ -1072,6 +1078,8 @@ model_list = list( | deepgram_models | elevenlabs_models | dashscope_models + | qwencloud_models + | qwen_ai_platform_models | moonshot_models | publicai_models | darkbloom_models @@ -1178,6 +1186,8 @@ def _build_models_by_provider() -> dict: "elevenlabs": elevenlabs_models, "heroku": heroku_models, "dashscope": dashscope_models, + "qwencloud": qwencloud_models, + "qwen_ai_platform": qwen_ai_platform_models, "modelscope": modelscope_models, "moonshot": moonshot_models, "publicai": publicai_models, @@ -2014,6 +2024,24 @@ if TYPE_CHECKING: from .llms.dashscope.rerank.transformation import ( DashScopeRerankConfig as DashScopeRerankConfig, ) + from .llms.dashscope.qwencloud import ( + QwenCloudChatConfig as QwenCloudChatConfig, + ) + from .llms.dashscope.qwencloud import ( + QwenCloudEmbeddingConfig as QwenCloudEmbeddingConfig, + ) + from .llms.dashscope.qwencloud import ( + QwenCloudRerankConfig as QwenCloudRerankConfig, + ) + from .llms.dashscope.qwen_ai_platform import ( + QwenAIPlatformChatConfig as QwenAIPlatformChatConfig, + ) + from .llms.dashscope.qwen_ai_platform import ( + QwenAIPlatformEmbeddingConfig as QwenAIPlatformEmbeddingConfig, + ) + from .llms.dashscope.qwen_ai_platform import ( + QwenAIPlatformRerankConfig as QwenAIPlatformRerankConfig, + ) from .llms.modelscope.chat.transformation import ( ModelScopeChatConfig as ModelScopeChatConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 1c833256598..e9199e1ec80 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -310,6 +310,8 @@ LLM_CONFIG_NAMES: Final = ( "GigaChatConfig", "GigaChatEmbeddingConfig", "DashScopeChatConfig", + "QwenCloudChatConfig", + "QwenAIPlatformChatConfig", "ModelScopeChatConfig", "MoonshotChatConfig", "DockerModelRunnerChatConfig", @@ -1172,6 +1174,14 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.dashscope.chat.transformation", "DashScopeChatConfig", ), + "QwenCloudChatConfig": ( + ".llms.dashscope.qwencloud", + "QwenCloudChatConfig", + ), + "QwenAIPlatformChatConfig": ( + ".llms.dashscope.qwen_ai_platform", + "QwenAIPlatformChatConfig", + ), "GDCGeminiConfig": ( ".llms.gdc.chat.transformation", "GDCGeminiConfig", diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index cefe6aae9ed..754815fce47 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -12,6 +12,7 @@ import hashlib import json import time import traceback +from collections.abc import Mapping from enum import Enum from typing import Any, Final @@ -506,7 +507,7 @@ class Cache: def _get_cache_logic( self, - cached_result: Any | None, + cached_result: object | None, max_age: float | None, ): """ @@ -538,8 +539,8 @@ class Cache: return cached_result @staticmethod - def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - cache_lookup_kwargs: Final[dict[str, Any]] = {} + def _get_safe_cache_lookup_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: + cache_lookup_kwargs: Final[dict[str, object]] = {} for prompt_kwarg in ("messages", "input"): if prompt_kwarg in kwargs: cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] @@ -552,7 +553,7 @@ class Cache: @staticmethod def _update_metadata_from_cache_lookup_kwargs( - original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any] + original_kwargs: Mapping[str, object], cache_lookup_kwargs: Mapping[str, object] ) -> None: original_metadata: Final = original_kwargs.get("metadata") cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata") diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 4898700c403..c5876e993d3 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -12,7 +12,7 @@ import ast import asyncio import json import os -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm from litellm._logging import print_verbose @@ -39,6 +39,12 @@ if TYPE_CHECKING: from litellm.router import Router +class _QdrantCollectionDetailsResponse(Protocol): + """The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object.""" + + def json(self) -> dict[str, object]: ... + + class QdrantSemanticCache(BaseCache): CACHE_KEY_FIELD_NAME = "litellm_cache_key" embedding_max_input_tokens: int | None = None @@ -115,15 +121,15 @@ class QdrantSemanticCache(BaseCache): raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}") if collection_exists.json()["result"]["exists"]: - collection_details = self.sync_client.get( + collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get( url=f"{self.qdrant_api_base}/collections/{self.collection_name}", headers=self.headers, ) - self.collection_info = collection_details.json() + self.collection_info: dict[str, object] = collection_details.json() print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}") self._ensure_cache_key_payload_index() else: - quantization_params: dict[str, Any] + quantization_params: dict[str, dict[str, object]] if quantization_config is None or quantization_config == "binary": quantization_params = { "binary": { @@ -214,7 +220,7 @@ class QdrantSemanticCache(BaseCache): resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router), ) - def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: + def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse: """Embed via the proxy Router when it serves the model, else direct.""" try: from litellm.proxy.proxy_server import llm_model_list, llm_router @@ -241,7 +247,7 @@ class QdrantSemanticCache(BaseCache): num_retries=0, ) - async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: + async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse: try: from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f1c80eaacbe..2b04a075114 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,7 @@ import time from collections.abc import Awaitable, Callable, Sequence from contextvars import ContextVar from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -58,6 +58,26 @@ else: Span = Any +class _AsyncRedisCommands(Protocol): + """Async redis commands this cache issues. + + redis-py's type stubs omit these methods on RedisCluster, so the union returned by + init_async_client() is untyped at every call site without this protocol. + """ + + def ping(self) -> Awaitable[bool]: ... + + def delete(self, *names: str) -> Awaitable[int]: ... + + def ttl(self, name: str) -> Awaitable[int]: ... + + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... + + def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... + + def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + + def _get_call_stack_info(num_frames: int = 2) -> str: """ Get the function names from the previous 1-2 functions in the call stack. @@ -429,6 +449,9 @@ class RedisCache(BaseCache): self.redis_async_client = redis_async_client return redis_async_client + def _async_commands(self) -> _AsyncRedisCommands: + return self.init_async_client() + def check_and_fix_namespace(self, key: str) -> str: """ Make sure each key starts with the given namespace @@ -1055,19 +1078,17 @@ class RedisCache(BaseCache): await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) self.redis_batch_writing_buffer = [] - def _get_cache_logic(self, cached_response: Any): + def _get_cache_logic(self, cached_response: bytes | str | None): """ Common 'get_cache_logic' across sync + async redis client implementations """ if cached_response is None: - return cached_response - # cached_response is in `b{} convert it to ModelResponse - cached_response = cached_response.decode("utf-8") # Convert bytes to string + return None + decoded: Final = cached_response.decode("utf-8") if isinstance(cached_response, bytes) else cached_response try: - cached_response = json.loads(cached_response) # Convert string to dictionary + return json.loads(decoded) except Exception: - cached_response = ast.literal_eval(cached_response) - return cached_response + return ast.literal_eval(decoded) def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): try: @@ -1314,8 +1335,7 @@ class RedisCache(BaseCache): raise e async def ping(self) -> bool: - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ping` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() print_verbose("Pinging Async Redis Cache") try: @@ -1349,8 +1369,7 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def delete_cache_keys(self, keys): - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() keys = [self.check_and_fix_namespace(key=key) for key in keys] # keys is a list, unpack it so it gets passed as individual elements to delete await _redis_client.delete(*keys) @@ -1415,8 +1434,7 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_delete_cache(self, key: str): - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) # keys is str return await _redis_client.delete(key) @@ -1523,8 +1541,7 @@ class RedisCache(BaseCache): Redis ref: https://redis.io/docs/latest/commands/ttl/ """ try: - # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl` - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) ttl: Final = await _redis_client.ttl(key) if ttl <= -1: # -1 means the key does not exist, -2 key does not exist @@ -1554,7 +1571,7 @@ class RedisCache(BaseCache): Returns: int: The length of the list after the push operation """ - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() try: @@ -1621,7 +1638,7 @@ class RedisCache(BaseCache): if len(rpush_list) == 0: return [] - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() try: @@ -1678,7 +1695,7 @@ class RedisCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> Any | list[Any]: - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") @@ -1810,7 +1827,7 @@ class RedisCache(BaseCache): if len(lpop_list) == 0: return [] - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final = self._async_commands() start_time: Final = time.time() try: diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 727c39c16ec..f494d6610a1 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -45,14 +45,14 @@ class ResponsesToCompletionBridgeHandler: return bool(stream) @staticmethod - def _is_preformatted_cached_chat_stream(result: Any) -> bool: + def _is_preformatted_cached_chat_stream(result: object) -> bool: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response" @staticmethod def _coerce_response_object( - response_obj: Any, + response_obj: object, hidden_params: dict | None, ) -> "ResponsesAPIResponse": if isinstance(response_obj, ResponsesAPIResponse): @@ -78,8 +78,8 @@ class ResponsesToCompletionBridgeHandler: for _ in stream_iter: pass - completed: Final = getattr(stream_iter, "completed_response", None) - response_obj: Final = getattr(completed, "response", None) if completed else None + completed: Final[object] = getattr(stream_iter, "completed_response", None) + response_obj: Final[object] = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") @@ -93,8 +93,8 @@ class ResponsesToCompletionBridgeHandler: async for _ in stream_iter: pass - completed: Final = getattr(stream_iter, "completed_response", None) - response_obj: Final = getattr(completed, "response", None) if completed else None + completed: Final[object] = getattr(stream_iter, "completed_response", None) + response_obj: Final[object] = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") @@ -157,7 +157,7 @@ class ResponsesToCompletionBridgeHandler: def completion( self, *args, **kwargs ) -> Union[ - Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", ]: diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 85fb0bc8dc6..7368de1e968 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -212,7 +212,8 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch LiteLLMCompletionResponsesConfig, ) - is_custom: Final = item.get("type") == "custom_tool_call" + item_type: Final[object] = item.get("type") + is_custom: Final = item_type == "custom_tool_call" arguments: Final = (item.get("input") if is_custom else item.get("arguments")) or "" name: Final = item.get("name") or ("custom_tool" if is_custom else "") function_chunk: Final = ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments) @@ -222,7 +223,7 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch function=function_chunk, index=index, ) - raw_provider_fields: Final = item.get("provider_specific_fields") + raw_provider_fields: Final[object] = item.get("provider_specific_fields") if isinstance(raw_provider_fields, dict): provider_specific_fields = raw_provider_fields elif raw_provider_fields and hasattr(raw_provider_fields, "__dict__"): @@ -507,7 +508,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _merge_responses_api_request_into_request_data( self, - request_data: dict[str, Any], + request_data: dict[str, object], responses_api_request: "ResponsesAPIOptionalRequestParams", instructions: str | None, ) -> None: diff --git a/litellm/constants.py b/litellm/constants.py index c482ab0e39a..1bd977dd9a9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -13,6 +13,12 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) +# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html +MAX_S3_OBJECT_KEY_BYTES: Final = 1024 +S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 +S3_PREFIX_DIGEST_CHARS: Final = 16 +# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against +MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) @@ -130,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_TOOL_LISTING_MAX_PAGES: Final = 1000 # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. @@ -630,6 +637,8 @@ LITELLM_CHAT_PROVIDERS: Final = [ "nscale", "nebius", "dashscope", + "qwencloud", + "qwen_ai_platform", "modelscope", "moonshot", "publicai", @@ -799,6 +808,7 @@ openai_compatible_endpoints: Final[list] = [ "inference.api.nscale.com/v1", "api.studio.nebius.ai/v1", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + "https://dashscope.aliyuncs.com/compatible-mode/v1", "https://api-inference.modelscope.cn/v1", "https://api.moonshot.ai/v1", "https://api.publicai.co/v1", @@ -872,6 +882,8 @@ openai_compatible_providers: Final[list] = [ "nscale", "nebius", "dashscope", + "qwencloud", + "qwen_ai_platform", "modelscope", "moonshot", "v0", @@ -902,6 +914,8 @@ openai_text_completion_compatible_providers: Final[list] = [ # providers that s "featherless_ai", "nebius", "dashscope", + "qwencloud", + "qwen_ai_platform", "modelscope", "moonshot", "publicai", @@ -1109,7 +1123,7 @@ nebius_models: Final[set] = set( ] ) -dashscope_models: Final[set] = set( +dashscope_models: Final[frozenset] = frozenset( [ "qwen-turbo", "qwen-plus", @@ -1124,6 +1138,10 @@ dashscope_models: Final[set] = set( ] ) +qwencloud_models: Final[frozenset] = frozenset(dashscope_models) + +qwen_ai_platform_models: Final[frozenset] = frozenset(dashscope_models) + nebius_embedding_models: Final[set] = set( [ "BAAI/bge-en-icl", @@ -1240,6 +1258,7 @@ BEDROCK_CONVERSE_MODELS: Final = [ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-fable-5-1", "anthropic.claude-fable-5", "anthropic.claude-sonnet-5", "anthropic.claude-opus-5", @@ -1571,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job" MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job" +USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job" PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job" SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report" SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3adc1c25dfd..b83e9b395a8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -641,12 +641,12 @@ def cost_per_token( return xai_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "lemonade": return lemonade_cost_per_token(model=model, usage=usage_block) - elif custom_llm_provider == "dashscope": + elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"): from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, ) - return dashscope_cost_per_token(model=model, usage=usage_block) + return dashscope_cost_per_token(model=model, usage=usage_block, custom_llm_provider=custom_llm_provider) elif custom_llm_provider == "azure_ai": return azure_ai_cost_per_token( model=model, @@ -1910,12 +1910,15 @@ def ocr_cost( if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: float | None = None - if model_info is not None: - ocr_cost_per_page = model_info.get("ocr_cost_per_page") + ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None + annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed - if pages_processed is None: + annotation_pages: Final = response.usage_info.pages_processed_annotation or 0 + has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0 + + if pages_processed is None and not has_billable_annotation_pages: if cost_per_credit is not None or ocr_cost_per_page is None: # Surface missing usage data instead of silently under-reporting # cost. The previous behavior raised ValueError; we now return 0.0 @@ -1931,7 +1934,7 @@ def ocr_cost( return 0.0, 0.0 raise ValueError("OCR response pages_processed is None") - if ocr_cost_per_page is None: + if ocr_cost_per_page is None and not has_billable_annotation_pages: # No per-page pricing configured. Either the model is on credit-based # pricing (and credits weren't returned, so the credit branch above did # not match) or the model has no OCR pricing entry at all. Surface a @@ -1947,8 +1950,9 @@ def ocr_cost( ) return 0.0, 0.0 - total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed - return total_ocr_processing_cost, 0.0 + ocr_pages_cost: Final = (ocr_cost_per_page or 0.0) * (pages_processed or 0) + annotation_pages_cost: Final = (annotation_rate or 0.0) * annotation_pages + return ocr_pages_cost + annotation_pages_cost, 0.0 def vector_store_search_cost( @@ -2268,6 +2272,10 @@ def batch_cost_calculator( return total_prompt_cost, total_completion_cost +def _attribute_value(obj: object, name: str) -> object: + return getattr(obj, name) + + def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]: field_names: Final = list(type(prompt_tokens_details).model_fields) if getattr(prompt_tokens_details, "cache_write_tokens", None) is None: @@ -2293,7 +2301,7 @@ class BaseTokenUsageProcessor: for usage in usage_objects: # Handle direct attributes by checking what exists in the model for attr in dir(usage): - if not attr.startswith("_") and not callable(getattr(usage, attr)): + if not attr.startswith("_") and not callable(_attribute_value(usage, attr)): current_val = getattr(combined, attr, 0) new_val = getattr(usage, attr, 0) if ( @@ -2313,7 +2321,7 @@ class BaseTokenUsageProcessor: if ( hasattr(usage.prompt_tokens_details, attr) and not attr.startswith("_") - and not callable(getattr(usage.prompt_tokens_details, attr)) + and not callable(_attribute_value(usage.prompt_tokens_details, attr)) ): current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0 new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0 @@ -2332,7 +2340,9 @@ class BaseTokenUsageProcessor: # Check what keys exist in the model's completion_tokens_details # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings for attr in type(usage.completion_tokens_details).model_fields: - if not attr.startswith("_") and not callable(getattr(usage.completion_tokens_details, attr)): + if not attr.startswith("_") and not callable( + _attribute_value(usage.completion_tokens_details, attr) + ): current_val = getattr(combined.completion_tokens_details, attr, 0) or 0 new_val = getattr(usage.completion_tokens_details, attr, 0) or 0 if isinstance(new_val, (int, float)): diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 9e949db625a..6c33621ec89 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -115,9 +115,11 @@ class SpeechToCompletionBridgeHandler: **request_data, ) + requested_response_format: Final = optional_params.get("response_format") if isinstance(result, ModelResponse): return self.transformation_handler.transform_response( model_response=result, + response_format=requested_response_format if isinstance(requested_response_format, str) else None, ) else: raise Exception(f"Unmapped response type. Got type: {type(result)}") diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 9b757ce86be..2ed140c0208 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -21,6 +21,8 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None: GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16" +GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm" +GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT}) class ChatAudioParam(TypedDict): @@ -29,6 +31,26 @@ class ChatAudioParam(TypedDict): class SpeechToCompletionBridgeTransformationHandler: + def _validate_response_format( + self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object] + ) -> None: + if not self._is_gemini_tts_model(model): + return + response_format: Final = optional_params.get("response_format") + if not isinstance(response_format, str) or response_format in GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: + return + from litellm.exceptions import BadRequestError + + supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS)) + raise BadRequestError( + message=( + f"Gemini TTS only produces raw PCM16 audio, so response_format='{response_format}'" + f" is not supported. Supported response formats: {supported}." + ), + model=model, + llm_provider=custom_llm_provider, + ) + def _chat_completion_params(self, optional_params: Mapping[str, object]) -> Mapping[str, object]: return MappingProxyType( { @@ -67,6 +89,7 @@ class SpeechToCompletionBridgeTransformationHandler: litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> dict: + self._validate_response_format(model, custom_llm_provider, optional_params) user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": input} return_kwargs: Final = { "model": model, @@ -125,7 +148,14 @@ class SpeechToCompletionBridgeTransformationHandler: """Check if the model is a Gemini TTS model that returns PCM16 data.""" return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower()) - def transform_response(self, model_response: "ModelResponse") -> "HttpxBinaryResponseContent": + def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]: + if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT: + return decoded_audio, "audio/pcm" + return self._convert_pcm16_to_wav(decoded_audio), "audio/wav" + + def transform_response( + self, model_response: "ModelResponse", response_format: str | None + ) -> "HttpxBinaryResponseContent": import base64 import httpx @@ -136,23 +166,17 @@ class SpeechToCompletionBridgeTransformationHandler: audio_part: Final = cast(Choices, model_response.choices[0]).message.audio if audio_part is None: raise ValueError("No audio part found in the response") - audio_content: Final = audio_part.data + decoded_audio: Final = base64.b64decode(audio_part.data) - # Decode base64 to get binary content - binary_data = base64.b64decode(audio_content) - - # Check if this is a Gemini TTS model that returns raw PCM16 data model: Final = getattr(model_response, "model", "") - headers: Final = {} - if self._is_gemini_tts_model(model): - # Convert PCM16 to WAV format for proper audio file playback - binary_data = self._convert_pcm16_to_wav(binary_data) - headers["Content-Type"] = "audio/wav" - else: - headers["Content-Type"] = "audio/mpeg" - - # Create an httpx.Response object - response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers) + content, content_type = ( + self._gemini_tts_response_body(decoded_audio, response_format) + if self._is_gemini_tts_model(model) + else (decoded_audio, "audio/mpeg") + ) + response: Final = httpx.Response( + status_code=200, content=content, headers=MappingProxyType({"Content-Type": content_type}) + ) binary_response: Final = HttpxBinaryResponseContent(response) binary_response.set_response_cost(_completion_response_cost(model_response)) return binary_response diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index f0a1bff8fdc..ea81e323da4 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -7,6 +7,7 @@ import base64 import os from collections.abc import Awaitable, Callable, Generator from datetime import timedelta +from functools import partial from importlib import metadata from typing import Any, Final, TypeVar @@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT +from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -603,17 +605,19 @@ class MCPClient: """ verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_tools_operation(session: ClientSession): - return await session.list_tools() - try: - result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count: Final = len(result.tools) - tool_names: Final = [tool.name for tool in result.tools] + # A per-server timeout above the global default extends the whole-walk deadline + listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT) + tools: Final = await self.run_with_session( + partial(list_tools_with_pagination, listing_deadline=listing_deadline), + quiet_on_error=raise_on_error, + ) + tool_count: Final = len(tools) + tool_names: Final = tuple(tool.name for tool in tools) verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) - return result.tools + return tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 30d50e2a74b..51d2139ef3b 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -1,14 +1,22 @@ import json from typing import Final, Literal +import anyio from mcp import ClientSession from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult +from mcp.types import PaginatedRequestParams from mcp.types import Tool as MCPTool from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_CLIENT_TIMEOUT, + MCP_TOOL_LISTING_MAX_PAGES, + MCP_TOOL_LISTING_TIMEOUT, +) from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages ) +async def list_tools_with_pagination( + session: ClientSession, listing_deadline: float | None = None +) -> list[MCPTool]: # mutable-ok: list return contract + """Collect tools from every tools/list page by following nextCursor. + + Stops and returns the tools collected so far when the upstream repeats a + cursor, the page cap is reached, or the whole-walk deadline expires, so a + buggy or slow upstream yields a partial catalog instead of an error. + listing_deadline overrides the default whole-walk deadline; callers with a + per-server timeout above the global default pass it through here. + """ + tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools + seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops + cursor: str | None = None # rebind-ok: advances to each page's nextCursor + # The per-request session read timeout restarts on every page, so a multi-page + # walk needs its own overall deadline. max() keeps the pre-pagination guarantee + # that a single page slower than the listing timeout but within the client + # timeout still succeeds. + effective_deadline: Final = ( + listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT) + ) + + with anyio.move_on_after(effective_deadline): + for _ in range(MCP_TOOL_LISTING_MAX_PAGES): + result = ( + await session.list_tools() + if cursor is None + else await session.list_tools(params=PaginatedRequestParams(cursor=cursor)) + ) + tools.extend(result.tools) + + next_cursor = getattr(result, "nextCursor", None) + if not isinstance(next_cursor, str) or not next_cursor: + return tools + if next_cursor in seen_cursors: + verbose_logger.warning( + "MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far", + len(tools), + ) + return tools + seen_cursors.add(next_cursor) + cursor = next_cursor + + verbose_logger.warning( + "MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far", + MCP_TOOL_LISTING_MAX_PAGES, + len(tools), + ) + return tools + + verbose_logger.warning( + "MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far", + effective_deadline, + len(tools), + ) + return tools + + async def load_mcp_tools( session: ClientSession, format: Literal["mcp", "openai"] = "mcp" ) -> list[MCPTool] | list[ChatCompletionToolParam]: @@ -103,10 +169,12 @@ async def load_mcp_tools( If format is set to "openai", the tools are converted to OpenAI API compatible tools. """ - tools: Final = await session.list_tools() + tools: Final = await list_tools_with_pagination(session) if format == "openai": - return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] - return tools.tools + return [ # mutable-ok: public API returns a list + transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools + ] + return tools ######################################################## diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 8ea19deef4c..6a698bb6018 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -722,7 +722,7 @@ class GoogleGenAIAdapter: ) for tool_call in tool_calls: - if not hasattr(tool_call, "function"): + if not hasattr(tool_call, "function") or isinstance(tool_call, ChatCompletionDeltaCustomToolCall): continue # 3. Use `index` as the primary key for accumulation diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index b5815bd3f7c..c1822e4720d 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -52,10 +52,10 @@ class GenerateContentSetupResult(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) model: str - request_body: dict[str, Any] + request_body: dict[str, object] custom_llm_provider: str generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None - generate_content_config_dict: dict[str, Any] + generate_content_config_dict: dict[str, object] native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj @@ -68,7 +68,7 @@ class GenerateContentHelper: @staticmethod def mock_generate_content_response( mock_response: str = "This is a mock response from Google GenAI generate_content.", - ) -> dict[str, Any]: + ) -> dict[str, object]: """Mock response for generate_content for testing purposes""" return { "text": mock_response, @@ -239,9 +239,9 @@ async def agenerate_content( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -307,9 +307,9 @@ def generate_content( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -397,9 +397,9 @@ async def agenerate_content_stream( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -492,9 +492,9 @@ def generate_content_stream( tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, diff --git a/litellm/images/main.py b/litellm/images/main.py index 1688087c2da..6a94e7c8df2 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -3,7 +3,7 @@ import contextvars import importlib from collections.abc import Coroutine from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload if TYPE_CHECKING: from litellm.images.utils import ImageEditRequestUtils @@ -151,7 +151,7 @@ def image_generation( *, aimg_generation: Literal[True], **kwargs, -) -> Coroutine[Any, Any, ImageResponse]: +) -> Coroutine[object, object, ImageResponse]: ... @@ -197,7 +197,7 @@ def image_generation( api_version: str | None = None, custom_llm_provider=None, **kwargs, -) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: +) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Maps the https://api.openai.com/v1/images/generations endpoint. @@ -386,6 +386,8 @@ def image_generation( litellm.LlmProviders.VERTEX_AI, litellm.LlmProviders.OPENROUTER, litellm.LlmProviders.DASHSCOPE, + litellm.LlmProviders.QWENCLOUD, + litellm.LlmProviders.QWEN_AI_PLATFORM, ): if image_generation_config is None: raise ValueError(f"image generation config is not supported for {custom_llm_provider}") @@ -723,14 +725,14 @@ def image_edit( user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: +) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Maps the image edit functionality, similar to OpenAI's images/edits endpoint. """ @@ -769,7 +771,7 @@ def image_edit( images: Final = image if isinstance(image, list) else ([image] if image is not None else []) headers_from_kwargs: Final = kwargs.get("headers") - merged_extra_headers: Final[dict[str, Any]] = {} + merged_extra_headers: Final[dict[str, object]] = {} if isinstance(headers_from_kwargs, dict): merged_extra_headers.update(headers_from_kwargs) if isinstance(extra_headers, dict): @@ -974,9 +976,9 @@ async def aimage_edit( user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1044,7 +1046,7 @@ async def aimage_edit( ) -def __getattr__(name: str) -> Any: +def __getattr__(name: str) -> type["ImageEditRequestUtils"]: """Lazy import handler for images.main module""" if name == "ImageEditRequestUtils": # Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94d734546be..748ef938cea 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from litellm.proxy.utils import PrismaClient from litellm.router import Router as _Router Router = _Router @@ -545,7 +546,6 @@ class SlackAlerting(CustomBatchLogger): # Get the appropriate budget alert type handler budget_alert_class: Final = get_budget_alert_type(type) _id: Final = budget_alert_class.get_id(user_info) - user_info_json: Final = user_info.model_dump(exclude_none=True) user_info_str: Final = self._get_user_info_str(user_info) event_message = budget_alert_class.get_event_message() @@ -575,7 +575,22 @@ class SlackAlerting(CustomBatchLogger): webhook_event = WebhookEvent( event=event, event_message=event_message, - **user_info_json, + spend=user_info.spend, + max_budget=user_info.max_budget, + soft_budget=user_info.soft_budget, + token=user_info.token, + customer_id=user_info.customer_id, + user_id=user_info.user_id, + team_id=user_info.team_id, + team_alias=user_info.team_alias, + organization_id=user_info.organization_id, + user_email=user_info.user_email, + key_alias=user_info.key_alias, + projected_exceeded_date=user_info.projected_exceeded_date, + projected_spend=user_info.projected_spend, + event_group=user_info.event_group, + alert_emails=user_info.alert_emails, + max_budget_alert_emails=user_info.max_budget_alert_emails, ) await self.send_alert( message=event_message + "\n\n" + user_info_str, @@ -657,7 +672,7 @@ class SlackAlerting(CustomBatchLogger): """ Create a standard message for a budget alert """ - _all_fields_as_dict: Final = user_info.model_dump(exclude_none=True) + _all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True) _all_fields_as_dict.pop("token") msg = "" for k, v in _all_fields_as_dict.items(): @@ -1006,7 +1021,7 @@ class SlackAlerting(CustomBatchLogger): except Exception: pass - async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any): + async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: object): base_model_from_user: Final = getattr(passed_model_info, "base_model", None) model_info = {} base_model = "" @@ -1930,6 +1945,69 @@ Model Info: except Exception as e: verbose_proxy_logger.exception("Error sending weekly spend report %s", e) + async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None: + """Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period.""" + if self.alerting is None or "slack" not in self.alerting: + return + + thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types + anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types + if not thresholds_enabled and not anomalies_enabled: + return + + if prisma_client is None: + from litellm.proxy.proxy_server import prisma_client as global_prisma_client + + prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client + if prisma_client is None: + return + + from litellm.integrations.SlackAlerting.user_spend_alerts import ( + evaluate_user_spend, + fetch_user_spend_rows, + ) + + try: + today: Final = datetime.datetime.now(datetime.timezone.utc).date() + rows: Final = await fetch_user_spend_rows( + prisma_client=prisma_client, + today=today, + baseline_days=self.alerting_args.spend_anomaly_baseline_days, + ) + all_events: Final = tuple( + event + for row in rows + for event in evaluate_user_spend( + row=row, + args=self.alerting_args, + today=today, + thresholds_enabled=thresholds_enabled, + anomalies_enabled=anomalies_enabled, + ) + ) + cached_flags: Final = await asyncio.gather( + *(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events) + ) + new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached) + for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies): + typed_events = tuple(event for event in new_events if event.alert_type == alert_type) + if not typed_events: + continue + await self.send_alert( + message="\n\n".join(event.message for event in typed_events), + level="High", + alert_type=alert_type, + alerting_metadata={}, # mutable-ok: send_alert takes a dict payload + ) + for event in typed_events: + await self.internal_usage_cache.async_set_cache( + key=event.cache_key, + value="SENT", + ttl=event.cache_ttl, + ) + except Exception as e: # noqa: BLE001 # background job must not crash the scheduler + verbose_proxy_logger.exception("Error sending user spend alerts: %s", e) + async def send_fallback_stats_from_prometheus(self): """ Helper to send fallback statistics from prometheus server -> to slack @@ -1973,7 +2051,7 @@ Model Info: try: message = f"`{event_name}`\n" - key_event_dict: Final = key_event.model_dump() + key_event_dict: Final[dict[str, object]] = key_event.model_dump() # Add Created by information first message += "*Action Done by:*\n" diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py new file mode 100644 index 00000000000..38794735c1b --- /dev/null +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -0,0 +1,139 @@ +"""Per-user daily/monthly spend threshold alerts and spend anomaly detection.""" + +import datetime +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Literal + +from pydantic import TypeAdapter + +from litellm.constants import HOURS_IN_A_DAY +from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60 +MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS + +USER_SPEND_QUERY: Final = """ +SELECT + user_id, + COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend, + COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend, + COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend +FROM "LiteLLM_DailyUserSpend" +WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL +GROUP BY user_id +HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0 +""" + + +@dataclass(frozen=True, slots=True) +class UserSpendRow: + user_id: str + daily_spend: float + monthly_spend: float + baseline_spend: float + + +@dataclass(frozen=True, slots=True) +class UserSpendAlertEvent: + kind: Literal["daily_threshold", "monthly_threshold", "anomaly"] + alert_type: AlertType + message: str + cache_key: str + cache_ttl: int + + +USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...]) + + +async def fetch_user_spend_rows( + prisma_client: "PrismaClient", + today: datetime.date, + baseline_days: int, +) -> tuple[UserSpendRow, ...]: + today_str: Final = today.strftime("%Y-%m-%d") + month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") + baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") + raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + return USER_SPEND_ROWS_ADAPTER.validate_python(raw) + + +def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: + threshold: Final = args.daily_spend_per_user_threshold + if threshold is None or row.daily_spend < threshold: + return None + return UserSpendAlertEvent( + kind="daily_threshold", + alert_type=AlertType.user_spend_thresholds, + message=( + f"User Daily Spend Threshold Crossed:\n" + f"User: `{row.user_id}`\n" + f"Spend Today: `${row.daily_spend:.2f}`\n" + f"Daily Threshold: `${threshold:.2f}`" + ), + cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}", + cache_ttl=DAY_SECONDS, + ) + + +def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None: + threshold: Final = args.monthly_spend_per_user_threshold + if threshold is None or row.monthly_spend < threshold: + return None + return UserSpendAlertEvent( + kind="monthly_threshold", + alert_type=AlertType.user_spend_thresholds, + message=( + f"User Monthly Spend Threshold Crossed:\n" + f"User: `{row.user_id}`\n" + f"Spend This Month: `${row.monthly_spend:.2f}`\n" + f"Monthly Threshold: `${threshold:.2f}`" + ), + cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}", + cache_ttl=MONTHLY_ALERT_TTL_SECONDS, + ) + + +def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: + if row.daily_spend < args.spend_anomaly_min_spend: + return None + baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days + if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg: + return None + return UserSpendAlertEvent( + kind="anomaly", + alert_type=AlertType.user_spend_anomalies, + message=( + f"User Spend Anomaly Detected:\n" + f"User: `{row.user_id}`\n" + f"Spend Today: `${row.daily_spend:.2f}`\n" + f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n" + f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average " + f"(minimum `${args.spend_anomaly_min_spend:.2f}`)" + ), + cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}", + cache_ttl=DAY_SECONDS, + ) + + +def evaluate_user_spend( + row: UserSpendRow, + args: SlackAlertingArgs, + today: datetime.date, + thresholds_enabled: bool, + anomalies_enabled: bool, +) -> tuple[UserSpendAlertEvent, ...]: + today_str: Final = today.strftime("%Y-%m-%d") + month_str: Final = today.strftime("%Y-%m") + threshold_events: Final = ( + ( + _daily_threshold_event(row=row, args=args, today_str=today_str), + _monthly_threshold_event(row=row, args=args, month_str=month_str), + ) + if thresholds_enabled + else () + ) + anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else () + return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None) diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index 71f4902bbe5..0c9e868c146 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -3,10 +3,12 @@ Arize Phoenix prompt manager that integrates with LiteLLM's prompt management sy Fetches prompt versions from Arize Phoenix and provides workspace-based access control. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, cast from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import ( @@ -20,6 +22,31 @@ from litellm.types.utils import StandardCallbackDynamicParams from .arize_phoenix_client import ArizePhoenixClient +class ArizePhoenixContentPart(TypedDict, total=False): + type: ReadOnly[str] + text: ReadOnly[str] + + +class ArizePhoenixTemplateMessage(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[Sequence[ArizePhoenixContentPart]] + + +class ArizePhoenixTemplateBody(TypedDict, total=False): + messages: ReadOnly[Sequence[ArizePhoenixTemplateMessage]] + + +class ArizePhoenixPromptMetadata(TypedDict): + model_name: ReadOnly[str | None] + model_provider: ReadOnly[str | None] + description: ReadOnly[str] + template_type: ReadOnly[str | None] + template_format: ReadOnly[str] + invocation_parameters: ReadOnly[Mapping[str, Mapping[str, object]]] + temperature: ReadOnly[float | None] + max_tokens: ReadOnly[int | None] + + class ArizePhoenixPromptTemplate: """ Represents a prompt template loaded from Arize Phoenix. @@ -28,10 +55,10 @@ class ArizePhoenixPromptTemplate: def __init__( self, template_id: str, - messages: list[dict[str, Any]], - metadata: dict[str, Any], + messages: Sequence[ArizePhoenixTemplateMessage], + metadata: ArizePhoenixPromptMetadata, model: str | None = None, - ): + ) -> None: self.template_id = template_id self.messages = messages self.metadata = metadata @@ -43,7 +70,7 @@ class ArizePhoenixPromptTemplate: self.description = metadata.get("description", "") self.template_format = metadata.get("template_format", "MUSTACHE") - def __repr__(self): + def __repr__(self) -> str: return f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')" @@ -109,7 +136,7 @@ class ArizePhoenixTemplateManager: def _parse_prompt_data(self, data: dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate: """Parse Arize Phoenix prompt data and extract messages and metadata.""" - template_data: Final = data.get("template", {}) + template_data: Final[ArizePhoenixTemplateBody] = data.get("template", {}) messages: Final = template_data.get("messages", []) # Extract invocation parameters @@ -129,7 +156,7 @@ class ArizePhoenixTemplateManager: break # Build metadata dictionary - metadata: Final = { + metadata: Final[ArizePhoenixPromptMetadata] = { "model_name": data.get("model_name"), "model_provider": data.get("model_provider"), "description": data.get("description", ""), @@ -146,7 +173,9 @@ class ArizePhoenixTemplateManager: metadata=metadata, ) - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> list[AllMessageValues]: + def render_template( + self, template_id: str, variables: Mapping[str, object] | None = None + ) -> list[AllMessageValues]: """Render a template with the given variables and return formatted messages.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -174,7 +203,9 @@ class ArizePhoenixTemplateManager: # Combine rendered content final_content = " ".join(rendered_content_parts) - rendered_messages.append({"role": role, "content": final_content}) + rendered_messages.append( + cast("AllMessageValues", {"role": role, "content": final_content}) # cast-ok: Phoenix roles are OpenAI + ) return rendered_messages @@ -243,8 +274,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, - ) -> tuple[list[AllMessageValues], dict[str, Any]]: + prompt_variables: Mapping[str, object] | None = None, + ) -> tuple[list[AllMessageValues], dict[str, object]]: """ Get a prompt template and render it with variables. @@ -263,7 +294,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): rendered_messages: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata - metadata: Final = { + metadata: Final[dict[str, object]] = { "model": template.model, "temperature": template.temperature, "max_tokens": template.max_tokens, @@ -271,7 +302,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement): # Add additional invocation parameters invocation_params: Final = template.invocation_parameters - provider_params = {} + provider_params: Mapping[str, object] = {} if "openai" in invocation_params: provider_params = invocation_params["openai"] @@ -289,12 +320,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: dict[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: dict[str, object] | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -335,9 +366,9 @@ class ArizePhoenixPromptManager(CustomPromptManagement): except Exception as e: # Log error but don't fail the call - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e) return messages, litellm_params def get_available_prompts(self) -> list[str]: @@ -393,7 +424,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Extract model from metadata (if specified) - template_model: Final = prompt_metadata.get("model") + raw_template_model: Final = prompt_metadata.get("model") + template_model: Final = raw_template_model if isinstance(raw_template_model, str) else None # Extract optional parameters from metadata optional_params: Final = {} diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index 6a03e3ee93c..ff34bd91e31 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -3,6 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system Fetches .prompt files from BitBucket repositories and provides team-based access control. """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from jinja2 import DictLoader, select_autoescape @@ -65,7 +66,7 @@ class BitBucketTemplateManager: def __init__( self, - bitbucket_config: dict[str, Any], + bitbucket_config: Mapping[str, object], prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config @@ -123,7 +124,7 @@ class BitBucketTemplateManager: template_content = content # Parse YAML frontmatter - metadata: dict[str, Any] = {} + metadata: dict[str, object] = {} if frontmatter_str: try: import yaml @@ -141,9 +142,9 @@ class BitBucketTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, object]: """Basic YAML parser for simple cases when PyYAML is not available.""" - result: Final[dict[str, Any]] = {} + result: Final[dict[str, object]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): @@ -162,7 +163,7 @@ class BitBucketTemplateManager: result[key] = value.strip("\"'") return result - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str: + def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str: """Render a template with the given variables.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") @@ -209,7 +210,7 @@ class BitBucketPromptManager(CustomPromptManagement): def __init__( self, - bitbucket_config: dict[str, Any], + bitbucket_config: Mapping[str, object], prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config @@ -234,7 +235,7 @@ class BitBucketPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, ) -> tuple[str, dict[str, Any]]: """ Get a prompt template and render it with variables. @@ -267,12 +268,12 @@ class BitBucketPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -316,9 +317,9 @@ class BitBucketPromptManager(CustomPromptManagement): except Exception as e: # Log error but don't fail the call - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e) return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]: @@ -384,14 +385,14 @@ class BitBucketPromptManager(CustomPromptManagement): def post_call_hook( self, user_id: str | None, - response: Any, + response: object, input_messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: Mapping[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> Any: + ) -> object: """ Post-call hook for any post-processing after the LLM call. """ diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index f0d4d67fc22..ffc8fe1c1f5 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -19,14 +19,29 @@ """Transform LiteLLM data to CloudZero AnyCost CBF format.""" from datetime import datetime -from typing import Any, Final +from typing import Final, SupportsFloat, SupportsIndex, SupportsInt import polars as pl +from typing_extensions import Buffer from ...types.integrations.cloudzero import CBFRecord from .cz_resource_names import CZEntityType, CZRNGenerator +def _as_int(value: object) -> int: + """The integer form of a spend table cell, computed the way :func:`int` computes it.""" + if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex)): + return int(value) + raise TypeError(f"int() argument must be a string or a number, not {type(value).__name__!r}") + + +def _as_float(value: object) -> float: + """The floating point form of a spend table cell, computed the way :func:`float` computes it.""" + if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)): + return float(value) + raise TypeError(f"float() argument must be a string or a number, not {type(value).__name__!r}") + + class CBFTransformer: """Transform LiteLLM usage data to CloudZero Billing Format (CBF).""" @@ -82,15 +97,15 @@ class CBFTransformer: return pl.DataFrame(cbf_data) - def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord: + def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord: """Create a single CBF record from LiteLLM daily spend row.""" # Parse date (daily spend tables use date strings like '2025-04-19') usage_date: Final = self._parse_date(row.get("date")) # Calculate total tokens - prompt_tokens: Final = int(row.get("prompt_tokens", 0)) - completion_tokens: Final = int(row.get("completion_tokens", 0)) + prompt_tokens: Final = _as_int(row.get("prompt_tokens", 0)) + completion_tokens: Final = _as_int(row.get("completion_tokens", 0)) total_tokens: Final = prompt_tokens + completion_tokens # Create CloudZero Resource Name (CZRN) as resource_id @@ -154,7 +169,7 @@ class CBFTransformer: "time/usage_start": ( usage_date.isoformat() if usage_date else None ), # Required: ISO-formatted UTC datetime - "cost/cost": float(row.get("spend", 0.0)), # Required: billed cost + "cost/cost": _as_float(row.get("spend", 0.0)), # Required: billed cost "resource/id": resource_id, # CZRN (CloudZero Resource Name) # Usage metrics for token consumption "usage/amount": total_tokens, # Numeric value of tokens consumed @@ -187,7 +202,7 @@ class CBFTransformer: return CBFRecord(cbf_record) - def _parse_date(self, date_str) -> datetime | None: + def _parse_date(self, date_str: object) -> datetime | None: """Parse date string from daily spend tables (e.g., '2025-04-19').""" if date_str is None: return None diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8dc6881d23e..372c9bf6b91 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,7 +1,9 @@ import contextvars +import copy import hashlib import os import secrets +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args @@ -38,6 +40,7 @@ except ImportError: if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation dc: Final = DualCache() @@ -227,13 +230,13 @@ class CustomGuardrail(CustomLogger): ) super().__init__(**kwargs) - def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str: + def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str: """Return a custom violation message if template is configured.""" if not self.violation_message_template: return default - format_context: Final[dict[str, Any]] = {"default_message": default} + format_context: Final[dict[str, object]] = {"default_message": default} if context: format_context.update(context) try: @@ -661,7 +664,7 @@ class CustomGuardrail(CustomLogger): value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails") return value if isinstance(value, list) else [] - def _is_valid_response_type(self, result: Any) -> bool: + def _is_valid_response_type(self, result: object) -> bool: """ Check if result is a valid LLMResponseTypes instance. @@ -722,7 +725,7 @@ class CustomGuardrail(CustomLogger): return None return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None: + def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -747,7 +750,7 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool: + def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: marker: Final = self._pre_call_marker() if marker is None: return False @@ -851,6 +854,69 @@ class CustomGuardrail(CustomLogger): return result + async def async_logging_hook( + self, + kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + result: object, + call_type: str, + ) -> tuple[dict, object]: # mutable-ok: CustomLogger.async_logging_hook contract + """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict.""" + from litellm.llms import get_guardrail_translation_mapping + + if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: + return kwargs, result + try: + translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))() + except ValueError: + verbose_logger.debug( + "Guardrail %s: no guardrail translation for call_type=%s, skipping logging_only scan", + self.guardrail_name, + call_type, + ) + return kwargs, result + litellm_params: Final = kwargs.get("litellm_params") or {} + scratch_metadata: Final = { + key: value + for key, value in (litellm_params.get("metadata") or {}).items() + if key != "standard_logging_guardrail_information" + } + try: + await self._scan_logged_call(kwargs, result, translation, scratch_metadata) + except Exception as e: + verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) + recorded: Final = scratch_metadata.get("standard_logging_guardrail_information") + standard_logging_object: Final = kwargs.get("standard_logging_object") + if not recorded or not isinstance(standard_logging_object, dict): + return kwargs, result + entries: Final = recorded if isinstance(recorded, list) else [recorded] + existing: Final = standard_logging_object.get("guardrail_information") or [] + return { + **kwargs, + "standard_logging_object": {**standard_logging_object, "guardrail_information": [*existing, *entries]}, + }, result + + async def _scan_logged_call( + self, + kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + result: object, + translation: "BaseTranslation", + scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata + ) -> None: + optional_params: Final = kwargs.get("optional_params") or {} + scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) + scratch_request: Final = { + "model": kwargs.get("model"), + "messages": scratch_input, + "input": scratch_input, + "tools": copy.deepcopy(optional_params.get("tools")), + "litellm_call_id": kwargs.get("litellm_call_id"), + "metadata": scratch_metadata, + } + await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) + await translation.process_output_response( + response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request + ) + def supports_scan_only_tool_results(self) -> bool: """Whether this guardrail can scan tool-result content. @@ -1170,7 +1236,7 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ # Convert None to empty dict to satisfy type requirements - guardrail_response: dict[str, Any] | str = {} if response is None else response + guardrail_response: dict[str, object] | str = {} if response is None else response # For apply_guardrail functions in custom_code_guardrail scenario, # simplify the logged response to "allow", "deny", or "mask" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index f177076d8fe..8f03e08f02d 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -31,6 +31,9 @@ if TYPE_CHECKING: from litellm.caching.caching import DualCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import ( MCPPostCallResponseObject, @@ -39,7 +42,7 @@ if TYPE_CHECKING: ) from litellm.types.router import PreRoutingHookResponse - Span = _Span | Any + Span = _Span else: Span = Any LiteLLMLoggingObj = Any @@ -268,7 +271,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> list[dict]: return healthy_deployments - async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict | None: """ Allow modifying the request just before it's sent to the deployment. @@ -344,9 +349,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_post_call_streaming_deployment_hook( self, request_data: dict, - response_chunk: Any, + response_chunk: object, call_type: CallTypes | None, - ) -> Any | None: + ) -> object | None: """ Allow modifying streaming chunks just before they're returned to the user. @@ -378,7 +383,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ def translate_completion_output_params_streaming( - self, completion_stream: Any + self, completion_stream: object ) -> AdapterCompletionStreamWrapper | None: """ Translates the streaming chunk, from the OpenAI format to the custom format. @@ -418,9 +423,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: object, request_headers: dict[str, str] | None = None, - litellm_call_info: dict[str, Any] | None = None, + litellm_call_info: dict[str, object] | None = None, ) -> dict[str, str] | None: """ Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers. @@ -471,11 +476,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: pass - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result @@ -581,7 +586,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_should_run_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -642,8 +647,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -711,8 +716,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, - anthropic_messages_provider_config: Any, + response: object, + anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None", anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -728,7 +733,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_post_agentic_loop_response_hook( self, - response: Any, + response: object, plan: AgenticLoopPlan, kwargs: dict, ) -> Any: @@ -767,7 +772,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_should_run_chat_completion_agentic_loop( self, - response: Any, + response: object, model: str, messages: list[dict], tools: list[dict] | None, @@ -785,12 +790,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, + response: object, optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, kwargs: dict, - ) -> Any: + ) -> object: """ Hook to execute chat completion agentic loop based on context from should_run hook. """ @@ -800,7 +805,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac tools: dict, model: str, messages: list[dict], - response: Any, + response: object, optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, @@ -1056,7 +1061,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def _redact_base64( self, - value: Any, + value: object, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, ) -> object: @@ -1079,7 +1084,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return value - def _should_keep_content(self, content: Any) -> bool: + def _should_keep_content(self, content: object) -> bool: """Return True if this content item should be retained.""" if not isinstance(content, dict): return True diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 04f1c6dff15..866076a3c49 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -20,10 +20,11 @@ import time import traceback from collections.abc import Sequence from datetime import datetime as datetimeObj -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx from httpx import Response +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -62,6 +63,18 @@ from litellm.types.utils import StandardLoggingPayload from ..additional_logging_utils import AdditionalLoggingUtils +if TYPE_CHECKING: + from fastapi import HTTPException + + from litellm.proxy._types import UserAPIKeyAuth + + +class _DatadogLoggingKwargs(TypedDict, total=False): + """The subset of logging ``kwargs`` that the Datadog payload builder reads.""" + + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + + # max number of logs DD API can accept @@ -87,6 +100,11 @@ def _resolve_dd_batch_size() -> int: return max(1, min(value, DD_MAX_BATCH_SIZE)) +def _span_attribute(span: object, name: str) -> object: + """Read an optional attribute off whatever span object the active tracer hands back.""" + return getattr(span, name, None) + + class DataDogLogger( CustomBatchLogger, AdditionalLoggingUtils, @@ -271,9 +289,9 @@ class DataDogLogger( self, request_data: dict, original_exception: Exception, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", traceback_str: str | None = None, - ) -> Any | None: + ) -> "HTTPException | None": """ Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog. @@ -297,7 +315,7 @@ class DataDogLogger( status_code = int(_code) # Use project-standard sanitized user context when running in proxy - user_context: dict[str, Any] = {} + user_context: dict[str, object] = {} try: from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, @@ -553,8 +571,8 @@ class DataDogLogger( def create_datadog_logging_payload( self, - kwargs: dict | Any, - response_obj: Any, + kwargs: _DatadogLoggingKwargs, + response_obj: object, start_time: datetime.datetime, end_time: datetime.datetime, ) -> DatadogPayload: @@ -562,8 +580,8 @@ class DataDogLogger( Helper function to create a datadog payload for logging Args: - kwargs (Union[dict, Any]): request kwargs - response_obj (Any): llm api response + kwargs: request kwargs, read for its standard logging object + response_obj: llm api response start_time (datetime.datetime): start time of request end_time (datetime.datetime): end time of request @@ -625,7 +643,7 @@ class DataDogLogger( self, payload: ServiceLoggerPayload, error: str | None = "", - parent_otel_span: Any | None = None, + parent_otel_span: object = None, start_time: datetimeObj | float | None = None, end_time: float | datetimeObj | None = None, event_metadata: dict | None = None, @@ -659,7 +677,7 @@ class DataDogLogger( self, payload: ServiceLoggerPayload, error: str | None = "", - parent_otel_span: Any | None = None, + parent_otel_span: object = None, start_time: datetimeObj | float | None = None, end_time: float | datetimeObj | None = None, event_metadata: dict | None = None, @@ -696,7 +714,7 @@ class DataDogLogger( def _create_v0_logging_payload( self, - kwargs: dict | Any, + kwargs: dict, response_obj: Any, start_time: datetime.datetime, end_time: datetime.datetime, @@ -810,11 +828,11 @@ class DataDogLogger( if current_span is None: return None - trace_id: Final = getattr(current_span, "trace_id", None) + trace_id: Final = _span_attribute(current_span, "trace_id") if trace_id is None: return None - span_id: Final = getattr(current_span, "span_id", None) + span_id: Final = _span_attribute(current_span, "span_id") trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)} if span_id is not None: trace_context["span_id"] = str(span_id) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 704f0323e95..5e116b7301a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -9,7 +9,9 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp import asyncio import json import os +from collections.abc import Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import Any, Final, Literal import httpx @@ -29,12 +31,16 @@ from litellm.integrations.datadog.datadog_mock_client import ( ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, handle_any_messages_to_chat_completion_str_messages_conversion, ) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( CallTypes, @@ -43,6 +49,189 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, ) +_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""} +_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024 + + +def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]: + """The value at `key` when it is a mapping, else an empty one.""" + value: Final = source.get(key) + return value if isinstance(value, dict) else _EMPTY_MAPPING + + +def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + content: Final = message.get("content") + if not isinstance(content, list): + return () + return tuple(block for block in content if isinstance(block, dict)) + + +def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str: + """ + Arguments as the object LLM Obs types them as, or the raw string when they are not one. + + Strings past the size bound ship unparsed: decoding multiplies memory on hostile compact + JSON, and the raw string is what the intake receives either way. + """ + if not isinstance(raw_arguments, str): + return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments) + if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS: + return raw_arguments + parsed: Final = safe_json_loads(raw_arguments) + return parsed if isinstance(parsed, dict) else raw_arguments + + +def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]: + """ + The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect. + + OpenAI puts them in `tool_calls` with the callee nested under `function` and `arguments` + serialized; Anthropic puts them in `content` as `tool_use` blocks with `input` already an + object. LLM Obs reads `name` / `arguments` / `tool_id` either way. + """ + raw_tool_calls: Final = message.get("tool_calls") + openai_calls: Final = tuple( + ToolCall( + name=function.get("name", ""), + arguments=_to_dd_arguments(function.get("arguments", "")), + tool_id=tool_call.get("id", ""), + type=tool_call.get("type", "function"), + ) + for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ()) + if isinstance(tool_call, dict) + for function in [_mapping_field(tool_call, "function")] + ) + anthropic_calls: Final = tuple( + ToolCall( + name=block.get("name", ""), + arguments=_to_dd_arguments(block.get("input") or {}), + tool_id=block.get("id", ""), + type="tool_use", + ) + for block in _content_blocks(message) + if block.get("type") == "tool_use" + ) + return openai_calls + anthropic_calls + + +def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]: + """ + The tool results a message carries, linked back to the call each answers. + + OpenAI models a result as a whole `role: "tool"` message keyed by `tool_call_id`; + Anthropic nests `tool_result` blocks inside a user message, keyed by `tool_use_id`. + """ + + def to_result(tool_id: str, result: object) -> ToolResult: + return ToolResult( + name=tool_call_names.get(tool_id, ""), + result=result if isinstance(result, str) else safe_dumps(result), + tool_id=tool_id, + type="function", + ) + + if message.get("role") == "tool": + return (to_result(str(message.get("tool_call_id", "")), message.get("content") or ""),) + return tuple( + to_result(str(block.get("tool_use_id", "")), block.get("content") or "") + for block in _content_blocks(message) + if block.get("type") == "tool_result" + ) + + +def _tool_call_names_by_id(messages: Sequence[object]) -> Mapping[str, str]: + """Ids to tool names for result linking; reads names structurally and parses nothing.""" + openai_pairs: Final = tuple( + (tool_call.get("id"), function.get("name", "")) + for message in messages + if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) + for tool_call in message["tool_calls"] + if isinstance(tool_call, dict) + for function in [_mapping_field(tool_call, "function")] + ) + anthropic_pairs: Final = tuple( + (block.get("id"), block.get("name", "")) + for message in messages + if isinstance(message, dict) + for block in _content_blocks(message) + if block.get("type") == "tool_use" + ) + return MappingProxyType({str(tool_id): str(name) for tool_id, name in openai_pairs + anthropic_pairs if tool_id}) + + +def _to_dd_message(message: object, tool_call_names: Mapping[str, str]) -> Message: + """ + Map one chat message onto LLM Obs' Message schema, adding fields and never destroying content. + + Content collapses to its text only when it has text; a content list with none (tool blocks, + images) rides along unchanged so nothing the caller logged is lost. Tool calls and results + move into the fields the LLM Obs Tools panel reads, from both the OpenAI and Anthropic shapes. + """ + if not isinstance(message, dict): + converted: Final = handle_any_messages_to_chat_completion_str_messages_conversion(message) + return converted[0] if converted else _EMPTY_MESSAGE + + text: Final = convert_content_list_to_str(message) # pyright: ignore[reportArgumentType] # caller-supplied dict + original_content: Final = message.get("content") + content: Final = ( + text if text or not isinstance(original_content, list) or not original_content else original_content + ) + reasoning: Final = message.get("reasoning_content") + tool_calls: Final = _to_dd_tool_calls(message) + tool_results: Final = _to_dd_tool_results(message, tool_call_names) + dd_message: Final[Message] = { + "role": message.get("role", ""), + "content": content, + **({"reasoning_content": reasoning} if reasoning is not None else {}), + **({"tool_calls": tool_calls} if tool_calls else {}), + **({"tool_results": tool_results} if tool_results else {}), + } + return dd_message + + +def _to_dd_messages(messages: object) -> tuple[Message, ...]: + """Map a whole conversation, resolving each tool result against the calls that precede it.""" + if messages is None: + return () + if not isinstance(messages, list): + return tuple(handle_any_messages_to_chat_completion_str_messages_conversion(messages)) + tool_call_names: Final = _tool_call_names_by_id(messages) + return tuple(_to_dd_message(message, tool_call_names) for message in messages) + + +def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None: + function: Final = entry.get("function") + declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry + name: Final = declared.get("name") + if not name: + return None + schema: Final = declared.get("parameters") or declared.get("input_schema") + description: Final = declared.get("description", "") + if not isinstance(schema, dict): + return ToolDefinition(name=name, description=description) + return ToolDefinition(name=name, description=description, schema=schema) + + +def _to_dd_tool_definitions(model_parameters: object) -> tuple[ToolDefinition, ...]: + """ + Map the request's declared tools onto LLM Obs' ToolDefinition schema. + + Handles the wrapped chat-completions shape and the bare shape the Anthropic and + Responses surfaces use, since both reach this logger through `model_parameters`. + """ + if not isinstance(model_parameters, dict): + return () + raw_tools: Final = model_parameters.get("tools") or model_parameters.get("functions") + if not isinstance(raw_tools, list): + return () + return tuple( + definition + for entry in raw_tools + if isinstance(entry, dict) + if (definition := _to_dd_tool_definition(entry)) is not None + ) + class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): @@ -221,12 +410,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): if standard_logging_payload is None: raise Exception("DataDogLLMObs: standard_logging_object is not set") - messages = standard_logging_payload["messages"] - messages = self._ensure_string_content(messages=messages) - metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) - input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages)) + input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"])) output_meta: Final = OutputMeta( messages=self._get_response_messages( standard_logging_payload=standard_logging_payload, @@ -240,22 +426,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): if isinstance(metadata, dict): metadata_parent_id = metadata.get("parent_id") - meta: Final = Meta( - kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id), - input=input_meta, - output=output_meta, - metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload), - error=error_info, - ) + tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters")) + span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id) + payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload) - # Calculate metrics (you may need to adjust these based on available data) - metrics: Final = LLMMetrics( - input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)), - output_tokens=float(standard_logging_payload.get("completion_tokens", 0)), - total_tokens=float(standard_logging_payload.get("total_tokens", 0)), - total_cost=float(standard_logging_payload.get("response_cost", 0)), - time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload), - ) + meta: Final[Meta] = { + "kind": span_kind, + "input": input_meta, + "output": output_meta, + "metadata": payload_metadata, + "error": error_info, + **({"tool_definitions": tool_definitions} if tool_definitions else {}), + } + + metrics: Final = self._assemble_metrics(standard_logging_payload) payload: Final[LLMObsPayload] = LLMObsPayload( parent_id=metadata_parent_id if metadata_parent_id else "undefined", @@ -313,6 +497,45 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info + def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics: + """ + Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from. + + Cache counts resolve through the same owners the savings dashboard uses, so every provider + spelling is covered, and `non_cached_input_tokens` subtracts BOTH cache categories because + litellm's normalized prompt count includes both (the invariant the cost calculator's custom + pricing helper documents). A zero residual on a fully cached request is real data and is + emitted; a zero read or write count is absence and is not. + """ + prompt_tokens: Final = float(standard_logging_payload.get("prompt_tokens", 0)) + completion_tokens: Final = float(standard_logging_payload.get("completion_tokens", 0)) + total_tokens: Final = float(standard_logging_payload.get("total_tokens", 0)) + total_cost: Final = float(standard_logging_payload.get("response_cost", 0)) + time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload) + + raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object") + usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None + cache_read: Final = float(extract_cache_read_tokens(usage_object)) + cache_write: Final = float(extract_cache_creation_tokens(usage_object)) + + metrics: Final[LLMMetrics] = { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": total_tokens, + "total_cost": total_cost, + "time_to_first_token": time_to_first_token, + **( + { + **({"cache_read_input_tokens": cache_read} if cache_read else {}), + **({"cache_write_input_tokens": cache_write} if cache_write else {}), + "non_cached_input_tokens": max(prompt_tokens - cache_read - cache_write, 0.0), + } + if cache_read or cache_write + else {} + ), + } + return metrics + def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float: """ Get the time to first token in seconds @@ -334,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): def _get_response_messages( self, standard_logging_payload: StandardLoggingPayload, call_type: str | None - ) -> list[Any]: + ) -> tuple[Message, ...]: """ Get the messages from the response object @@ -343,7 +566,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): response_obj = standard_logging_payload.get("response") if response_obj is None: - return [] + return () # edge case: handle response_obj is a string representation of a dict if isinstance(response_obj, str): @@ -356,7 +579,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # fallback to json parsing response_obj = json.loads(str(response_obj)) except json.JSONDecodeError: - return [] + return () if call_type in [ CallTypes.completion.value, @@ -374,12 +597,12 @@ class DataDogLLMObsLogger(CustomBatchLogger): if isinstance(response_obj, dict) and "choices" in response_obj: choices: Final = response_obj["choices"] if choices and len(choices) > 0 and "message" in choices[0]: - return [choices[0]["message"]] - return [] + return _to_dd_messages([choices[0]["message"]]) + return () except (KeyError, IndexError, TypeError): # In case of any error accessing the response structure, return empty list - return [] - return [] + return () + return () def _get_datadog_span_kind( self, call_type: str | None, parent_id: str | None = None @@ -484,22 +707,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]: - if messages is None: - return [] - if isinstance(messages, str): - return [messages] - elif isinstance(messages, list): - return [message for message in messages] - elif isinstance(messages, dict): - return [str(messages.get("content", ""))] - return [] - - def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: + def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload """ - _metadata: Final[dict[str, Any]] = { + _metadata: Final[dict[str, object]] = { "model_name": standard_logging_payload.get("model", "unknown"), "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"), "id": standard_logging_payload.get("id", "unknown"), @@ -523,10 +735,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): spend_metrics: Final = self._get_spend_metrics(standard_logging_payload) _metadata.update({"spend_metrics": dict(spend_metrics)}) - ## extract tool calls and add to metadata - tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload) - _metadata.update(tool_call_metadata) - _standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {} _metadata.update(_standard_logging_metadata) return _metadata @@ -646,107 +854,3 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at) return spend_metrics - - def _process_input_messages_preserving_tool_calls(self, messages: list[Any]) -> list[dict[str, Any]]: - """ - Process input messages while preserving tool_calls and tool message types. - - This bypasses the lossy string conversion when tool calls are present, - allowing complex nested tool_calls objects to be preserved for Datadog. - """ - processed: Final = [] - for msg in messages: - if isinstance(msg, dict): - # Preserve messages with tool_calls or tool role as-is - if "tool_calls" in msg or msg.get("role") == "tool": - processed.append(msg) - else: - # For regular messages, still apply string conversion - converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) - processed.extend(converted) - else: - # For non-dict messages, apply string conversion - converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) - processed.extend(converted) - return processed - - @staticmethod - def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, Any]: - """ - Extract tool call information into key-value pairs for Datadog metadata. - - Similar to OpenTelemetry's implementation but adapted for Datadog's format. - """ - kv_pairs: Final[dict[str, Any]] = {} - for idx, tool_call in enumerate(tool_calls): - try: - # Extract tool call ID - tool_id = tool_call.get("id") - if tool_id: - kv_pairs[f"tool_calls.{idx}.id"] = tool_id - - # Extract tool call type - tool_type = tool_call.get("type") - if tool_type: - kv_pairs[f"tool_calls.{idx}.type"] = tool_type - - # Extract function information - function = tool_call.get("function") - if function: - function_name = function.get("name") - if function_name: - kv_pairs[f"tool_calls.{idx}.function.name"] = function_name - - function_arguments = function.get("arguments") - if function_arguments: - # Store arguments as JSON string for Datadog - if isinstance(function_arguments, str): - kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments - else: - import json - - kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments) - except (KeyError, TypeError, ValueError) as e: - verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e) - continue - - return kv_pairs - - def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: - """ - Extract tool call information from both input messages and response for Datadog metadata. - """ - tool_call_metadata: Final[dict[str, Any]] = {} - - try: - # Extract tool calls from input messages - messages: Final = standard_logging_payload.get("messages", []) - if messages and isinstance(messages, list): - for message in messages: - if isinstance(message, dict) and "tool_calls" in message: - tool_calls = message.get("tool_calls") - if tool_calls: - input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls) - # Prefix with "input_" to distinguish from response tool calls - for key, value in input_tool_calls_kv.items(): - tool_call_metadata[f"input_{key}"] = value - - # Extract tool calls from response - response_obj: Final = standard_logging_payload.get("response") - if response_obj and isinstance(response_obj, dict): - choices: Final = response_obj.get("choices", []) - for choice in choices: - if isinstance(choice, dict): - message = choice.get("message") - if message and isinstance(message, dict): - tool_calls = message.get("tool_calls") - if tool_calls: - response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls) - # Prefix with "output_" to distinguish from input tool calls - for key, value in response_tool_calls_kv.items(): - tool_call_metadata[f"output_{key}"] = value - - except Exception as e: - verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e) - - return tool_call_metadata diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index fd0b17ba746..9c82ff7c5ba 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -3,12 +3,21 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d """ import re +from collections.abc import Mapping from pathlib import Path from typing import Any, Final import yaml from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import NotRequired, ReadOnly, TypedDict + + +class _PromptFileJson(TypedDict): + """JSON form of a .prompt file: rendered template text plus its frontmatter.""" + + content: ReadOnly[NotRequired[str]] + metadata: ReadOnly[NotRequired[dict[str, object]]] def strip_version_suffix(prompt_id: str) -> str | None: @@ -167,7 +176,7 @@ class PromptManager: template_id=prompt_id, ) - def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]: + def _parse_frontmatter(self, content: str) -> tuple[dict[str, object], str]: """Parse YAML frontmatter from prompt content.""" # Match YAML frontmatter between --- delimiters frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$" @@ -178,7 +187,7 @@ class PromptManager: template_content = match.group(2) try: - frontmatter = yaml.safe_load(frontmatter_yaml) or {} + frontmatter: dict[str, object] = yaml.safe_load(frontmatter_yaml) or {} except yaml.YAMLError as e: raise ValueError(f"Invalid YAML frontmatter: {e}") else: @@ -191,7 +200,7 @@ class PromptManager: def render( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, version: int | None = None, ) -> str: """ @@ -231,7 +240,7 @@ class PromptManager: except Exception as e: raise ValueError(f"Error rendering template '{prompt_id}': {e}") - def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None: + def _validate_input(self, variables: Mapping[str, object], schema: Mapping[str, str]) -> None: """Basic validation of input variables against schema.""" for field_name, field_type in schema.items(): if field_name in variables: @@ -291,7 +300,7 @@ class PromptManager: """Get a list of all available prompt IDs.""" return list(self.prompts.keys()) - def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None: + def get_prompt_metadata(self, prompt_id: str) -> dict[str, object] | None: """Get metadata for a specific prompt.""" template: Final = self.prompts.get(prompt_id) return template.metadata if template else None @@ -302,12 +311,12 @@ class PromptManager: if self.prompt_directory: self._load_prompts() - def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None: + def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, object] | None = None) -> None: """Add a prompt template programmatically.""" template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id) self.prompts[prompt_id] = template - def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]: + def prompt_file_to_json(self, file_path: str | Path) -> _PromptFileJson: """Convert a .prompt file to JSON format. Args: @@ -324,7 +333,7 @@ class PromptManager: return {"content": template_content.strip(), "metadata": frontmatter} - def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str: + def json_to_prompt_file(self, prompt_data: _PromptFileJson) -> str: """Convert JSON prompt data to .prompt file format. Args: diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 23727801a6f..b27618993a3 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -6,10 +6,11 @@ import re import uuid from collections.abc import Mapping, Sequence from datetime import datetime, timezone, tzinfo -from typing import Any, Final, TypedDict, cast +from typing import Any, Final, Protocol, cast import httpx from pydantic import BaseModel, Field +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -35,6 +36,34 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai" GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000 +class _GalileoLoginBody(TypedDict): + """Decoded body of the Galileo login response.""" + + access_token: ReadOnly[str] + + +class _GalileoLoginResponse(Protocol): + """The login call's HTTP response, read for the access token it carries.""" + + def json(self) -> _GalileoLoginBody: ... + + +class _JsonResponse(Protocol): + """An HTTP response read only for whatever JSON body it decodes to.""" + + def json(self) -> object: ... + + +def _login_access_token(response: _GalileoLoginResponse) -> str: + """Read the bearer token out of a Galileo login response body.""" + return response.json()["access_token"] + + +def _decoded_body(response: _JsonResponse) -> object: + """Decode a response body without asserting anything about its shape.""" + return response.json() + + class GalileoStandardLoggingFields(TypedDict, total=False): call_type: str model: str @@ -156,7 +185,7 @@ class GalileoObserve(CustomLogger): }, ) galileo_login_response.raise_for_status() - access_token: Final = galileo_login_response.json()["access_token"] + access_token: Final = _login_access_token(galileo_login_response) self.headers = { "accept": "application/json", "Content-Type": "application/json", @@ -421,7 +450,7 @@ class GalileoObserve(CustomLogger): try: verbose_logger.debug( "Galileo Logger HTTP error response json: %s", - response.json(), + _decoded_body(response), ) except Exception: pass diff --git a/litellm/integrations/gitlab/gitlab_client.py b/litellm/integrations/gitlab/gitlab_client.py index 0690ccc8c15..813a2ef2821 100644 --- a/litellm/integrations/gitlab/gitlab_client.py +++ b/litellm/integrations/gitlab/gitlab_client.py @@ -4,12 +4,80 @@ Now supports selecting a tag via `config["tag"]`; falls back to branch ("main"). """ import base64 -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Protocol, TypedDict from urllib.parse import quote +from typing_extensions import ReadOnly + from litellm.llms.custom_httpx.http_handler import HTTPHandler +class GitLabFilePayload(TypedDict, total=False): + """A repository-files API entry.""" + + content: ReadOnly[str] + encoding: ReadOnly[str] + + +class GitLabTreeEntry(TypedDict, total=False): + """A repository-tree API entry.""" + + path: ReadOnly[str] + type: ReadOnly[str] + + +class GitLabBranch(TypedDict, total=False): + """A repository-branches API entry.""" + + name: ReadOnly[str] + type: ReadOnly[str] + + +class GitLabFileMetadata(TypedDict): + """The response headers a raw file request exposes as metadata.""" + + content_type: ReadOnly[str | None] + content_length: ReadOnly[str | None] + last_modified: ReadOnly[str | None] + + +class _FileJsonResponse(Protocol): + def json(self) -> GitLabFilePayload: ... + + +class _TreeJsonResponse(Protocol): + def json(self) -> Sequence[GitLabTreeEntry] | None: ... + + +class _ProjectJsonResponse(Protocol): + def json(self) -> Mapping[str, object]: ... + + +class _BranchesJsonResponse(Protocol): + def json(self) -> Sequence[GitLabBranch] | None: ... + + +def _file_payload(resp: _FileJsonResponse) -> GitLabFilePayload: + """The JSON body of a repository-files response.""" + return resp.json() + + +def _tree_entries(resp: _TreeJsonResponse) -> Sequence[GitLabTreeEntry]: + """The entries of a repository-tree response.""" + return resp.json() or [] + + +def _project_info(resp: _ProjectJsonResponse) -> Mapping[str, object]: + """The JSON body of a project response.""" + return resp.json() + + +def _branch_entries(resp: _BranchesJsonResponse) -> Sequence[GitLabBranch] | None: + """The JSON body of a repository-branches response.""" + return resp.json() + + class GitLabClient: """ Client for interacting with the GitLab API to fetch files. @@ -42,12 +110,12 @@ class GitLabClient: self.project: str | int = project self.access_token: str = str(access_token) - self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth' + self.auth_method: str = config.get("auth_method", "token") # 'token' or 'oauth' self.branch = config.get("branch", None) if not self.branch: self.branch = "main" self.tag = config.get("tag") - self.base_url = config.get("base_url", "https://gitlab.com/api/v4") + self.base_url: str = config.get("base_url", "https://gitlab.com/api/v4") if not all([self.project, self.access_token]): raise ValueError("project and access_token are required") @@ -159,7 +227,7 @@ class GitLabClient: if resp.status_code == 404: return None resp.raise_for_status() - data: Final = resp.json() + data: Final = _file_payload(resp) content: Final = data.get("content") encoding: Final = data.get("encoding", "") if content and encoding == "base64": @@ -208,7 +276,7 @@ class GitLabClient: return [] resp.raise_for_status() - data: Final = resp.json() or [] + data: Final = _tree_entries(resp) files: Final[list[str]] = [] for item in data: if item.get("type") == "blob": @@ -229,13 +297,13 @@ class GitLabClient: raise Exception("Authentication failed. Check your GitLab token and auth_method.") raise Exception(f"Failed to list files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """Get information about the project/repository.""" url: Final = f"{self.base_url}/projects/{self._project_enc}" try: resp: Final = self.http_handler.get(url, headers=self.headers) resp.raise_for_status() - return resp.json() + return _project_info(resp) except Exception as e: raise Exception(f"Failed to get repository info: {e}") @@ -247,18 +315,18 @@ class GitLabClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> list[GitLabBranch]: """Get list of branches in the repository.""" url: Final = f"{self.base_url}/projects/{self._project_enc}/repository/branches" try: resp: Final = self.http_handler.get(url, headers=self.headers) resp.raise_for_status() - data: Final = resp.json() + data: Final = _branch_entries(resp) return data if isinstance(data, list) else [] except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> GitLabFileMetadata | None: """ Get minimal metadata about a file via RAW endpoint headers at a given ref. diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index c41d9dd240f..d4602176650 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -2,10 +2,12 @@ GitLab prompt manager with configurable prompts folder. """ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, TypeVar from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import ReadOnly, TypedDict from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -24,6 +26,19 @@ from litellm.types.utils import StandardCallbackDynamicParams GITLAB_PREFIX: Final = "gitlab::" +_ResponseT = TypeVar("_ResponseT") + + +class GitLabCachedPrompt(TypedDict): + id: ReadOnly[str] + path: ReadOnly[str] + content: ReadOnly[str] + metadata: ReadOnly[Mapping[str, object]] + model: ReadOnly[str | None] + temperature: ReadOnly[float | None] + max_tokens: ReadOnly[int | None] + optional_params: ReadOnly[Mapping[str, object]] + def encode_prompt_id(raw_id: str) -> str: """Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'""" @@ -206,7 +221,7 @@ class GitLabTemplateManager: result[key] = value.strip("\"'") return result - def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str: + def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str: if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") template: Final = self.prompts[template_id] @@ -313,7 +328,7 @@ class GitLabPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, *, ref: str | None = None, ) -> tuple[str, dict[str, Any]]: @@ -338,13 +353,13 @@ class GitLabPromptManager(CustomPromptManagement): self, user_id: str | None, messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: dict[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, prompt_version: str | None = None, **kwargs, - ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: + ) -> tuple[list[AllMessageValues], dict[str, object] | None]: if not prompt_id: return messages, litellm_params try: @@ -377,9 +392,9 @@ class GitLabPromptManager(CustomPromptManagement): return final_messages, litellm_params except Exception as e: - import litellm + from litellm._logging import verbose_proxy_logger - litellm._logging.verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e) + verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e) return messages, litellm_params def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]: @@ -435,14 +450,14 @@ class GitLabPromptManager(CustomPromptManagement): def post_call_hook( self, user_id: str | None, - response: Any, + response: _ResponseT, input_messages: list[AllMessageValues], - function_call: dict[str, Any] | str | None = None, - litellm_params: dict[str, Any] | None = None, + function_call: Mapping[str, object] | str | None = None, + litellm_params: Mapping[str, object] | None = None, prompt_id: str | None = None, - prompt_variables: dict[str, Any] | None = None, + prompt_variables: Mapping[str, object] | None = None, **kwargs, - ) -> Any: + ) -> _ResponseT: return response def get_available_prompts(self) -> list[str]: @@ -498,7 +513,7 @@ class GitLabPromptManager(CustomPromptManagement): messages: Final = self._parse_prompt_to_messages(rendered_prompt) template_model: Final = prompt_metadata.get("model") - optional_params: Final[dict[str, Any]] = {} + optional_params: Final[dict[str, object]] = {} for param in [ "temperature", "max_tokens", @@ -658,14 +673,14 @@ class GitLabPromptCache: self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager # In-memory stores - self._by_file: dict[str, dict[str, Any]] = {} - self._by_id: dict[str, dict[str, Any]] = {} + self._by_file: dict[str, GitLabCachedPrompt] = {} + self._by_id: dict[str, GitLabCachedPrompt] = {} # ------------------------- # Public API # ------------------------- - def load_all(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]: + def load_all(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]: """ Scan GitLab for all .prompt files under prompts_path, load and parse each, and return the mapping of repo file path -> JSON-like dict. @@ -695,7 +710,7 @@ class GitLabPromptCache: return self._by_id - def reload(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]: + def reload(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]: """Clear the cache and re-load from GitLab.""" self._by_file.clear() self._by_id.clear() @@ -709,11 +724,11 @@ class GitLabPromptCache: """Return the template IDs (relative to prompts_path, without extension) currently cached.""" return list(self._by_id.keys()) - def get_by_file(self, file_path: str) -> dict[str, Any] | None: + def get_by_file(self, file_path: str) -> GitLabCachedPrompt | None: """Get a cached prompt JSON by repo file path.""" return self._by_file.get(file_path) - def get_by_id(self, prompt_id: str) -> dict[str, Any] | None: + def get_by_id(self, prompt_id: str) -> GitLabCachedPrompt | None: """Get a cached prompt JSON by prompt ID (relative to prompts_path).""" if prompt_id in self._by_id: return self._by_id[prompt_id] @@ -728,7 +743,7 @@ class GitLabPromptCache: # Internals # ------------------------- - def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> dict[str, Any]: + def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> GitLabCachedPrompt: """ Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize. """ diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 296c2b5714e..9576eabaa34 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -89,7 +89,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): - prompt_tokens_details: Final = getattr(usage_obj, "prompt_tokens_details", None) + prompt_tokens_details: Final[object] = getattr(usage_obj, "prompt_tokens_details", None) if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"): cached_tokens: Final = getattr(prompt_tokens_details, "cached_tokens", None) if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0: @@ -623,9 +623,16 @@ class LangFuseLogger: ) # Apply custom masking function if provided - if masking_function is not None and callable(masking_function): - input = self._apply_masking_function(input, masking_function) - output = self._apply_masking_function(output, masking_function) + masked_input: Final[object] = ( + self._apply_masking_function(input, masking_function) + if masking_function is not None and callable(masking_function) + else input + ) + masked_output: Final[object] = ( + self._apply_masking_function(output, masking_function) + if masking_function is not None and callable(masking_function) + else output + ) clean_metadata = redact_user_api_key_info(metadata=clean_metadata) @@ -651,15 +658,15 @@ class LangFuseLogger: # Special keys that are found in the function arguments and not the metadata if "input" in update_trace_keys: - trace_params["input"] = input if not mask_input else "redacted-by-litellm" + trace_params["input"] = masked_input if not mask_input else "redacted-by-litellm" if "output" in update_trace_keys: - trace_params["output"] = output if not mask_output else "redacted-by-litellm" + trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" else: # don't overwrite an existing trace trace_params = { "id": trace_id, "name": trace_name, "session_id": session_id, - "input": input if not mask_input else "redacted-by-litellm", + "input": masked_input if not mask_input else "redacted-by-litellm", "version": clean_metadata.pop( "trace_version", clean_metadata.get("version", None) ), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence @@ -669,9 +676,9 @@ class LangFuseLogger: trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None) if level == "ERROR": - trace_params["status_message"] = output + trace_params["status_message"] = masked_output else: - trace_params["output"] = output if not mask_output else "redacted-by-litellm" + trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm" if debug is True or (isinstance(debug, str) and debug.lower() == "true"): debug_metadata: Final = { @@ -708,7 +715,7 @@ class LangFuseLogger: ("aws_region_name", aws_region_name, bool(aws_region_name)), ("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs), ) - enrichments: Final[Mapping[str, Any]] = { + enrichments: Final[Mapping[str, object]] = { key: value for key, value, include in candidate_enrichments if include } @@ -802,8 +809,8 @@ class LangFuseLogger: "end_time": end_time, "model": model_name, "model_parameters": optional_params, - "input": input if not mask_input else "redacted-by-litellm", - "output": output if not mask_output else "redacted-by-litellm", + "input": masked_input if not mask_input else "redacted-by-litellm", + "output": masked_output if not mask_output else "redacted-by-litellm", "usage": usage, "usage_details": usage_details, "metadata": { @@ -825,8 +832,8 @@ class LangFuseLogger: prompt_management_metadata=prompt_management_metadata, langfuse_client=self.Langfuse, ) - if output is not None and isinstance(output, str) and level == "ERROR": - generation_params["status_message"] = output + if masked_output is not None and isinstance(masked_output, str) and level == "ERROR": + generation_params["status_message"] = masked_output if self._supports_completion_start_time(): generation_params["completion_start_time"] = kwargs.get("completion_start_time", None) @@ -935,7 +942,7 @@ class LangFuseLogger: return Version(self.langfuse_sdk_version) >= Version("2.7.3") @staticmethod - def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any: + def _apply_masking_function(data: object, masking_function: Callable[[object], object]) -> object: """ Apply a masking function to data, handling different data types. @@ -1049,7 +1056,7 @@ def _add_prompt_to_generation_params( generation_params: dict, clean_metadata: dict, prompt_management_metadata: StandardLoggingPromptManagementMetadata | None, - langfuse_client: Any, + langfuse_client: object, ) -> dict: from langfuse import Langfuse from langfuse.model import ( diff --git a/litellm/integrations/opik/opik.py b/litellm/integrations/opik/opik.py index fae93f03d1e..ce47d7fe27a 100644 --- a/litellm/integrations/opik/opik.py +++ b/litellm/integrations/opik/opik.py @@ -4,9 +4,12 @@ Opik Logger that logs LLM events to an Opik server import asyncio import traceback +from collections.abc import Mapping from datetime import datetime from typing import Any, Final +from typing_extensions import ReadOnly, TypedDict, Unpack + from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.llms.custom_httpx.http_handler import ( @@ -23,7 +26,7 @@ except Exception: opik_client = None -def _should_skip_event(kwargs: dict[str, Any]) -> bool: +def _should_skip_event(kwargs: Mapping[str, object]) -> bool: """Check if event should be skipped due to missing standard_logging_object.""" if kwargs.get("standard_logging_object") is None: verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found") @@ -31,12 +34,24 @@ def _should_skip_event(kwargs: dict[str, Any]) -> bool: return False +class _OpikLoggerKwargs(TypedDict, total=False): + """Constructor options accepted by ``OpikLogger``.""" + + project_name: ReadOnly[str | None] + url: ReadOnly[str | None] + api_key: ReadOnly[str | None] + workspace: ReadOnly[str | None] + batch_size: ReadOnly[int | None] + flush_interval: ReadOnly[int | None] + max_queue_size: ReadOnly[int | None] + + class OpikLogger(CustomBatchLogger): """ Opik Logger for logging events to an Opik Server """ - def __init__(self, **kwargs: Any) -> None: + def __init__(self, **kwargs: Unpack[_OpikLoggerKwargs]) -> None: self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.sync_httpx_client = _get_httpx_client() @@ -95,7 +110,7 @@ class OpikLogger(CustomBatchLogger): async def async_log_success_event( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], response_obj: Any, start_time: datetime, end_time: datetime, @@ -163,7 +178,7 @@ class OpikLogger(CustomBatchLogger): except Exception as e: verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc()) - def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: + def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None: try: response: Final = self.sync_httpx_client.post( url=url, @@ -178,7 +193,7 @@ class OpikLogger(CustomBatchLogger): def log_success_event( self, - kwargs: dict[str, Any], + kwargs: dict[str, object], response_obj: Any, start_time: datetime, end_time: datetime, @@ -247,7 +262,7 @@ class OpikLogger(CustomBatchLogger): except Exception as e: verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc()) - async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None: + async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None: try: response: Final = await self.async_httpx_client.post( url=url, diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index 92a7eca7f3e..4dd3d40fae3 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -1,6 +1,7 @@ """Data extraction functions for Opik payload building.""" import json +from collections.abc import Mapping from typing import Any, Final from litellm import _logging @@ -35,8 +36,8 @@ def normalize_provider_name(provider: str | None) -> str | None: def extract_opik_metadata( - litellm_metadata: dict[str, Any], - standard_logging_metadata: dict[str, Any], + litellm_metadata: Mapping[str, Any], + standard_logging_metadata: Mapping[str, Any], ) -> dict[str, Any]: """ Merge Opik metadata from three sources in increasing priority order: @@ -97,7 +98,7 @@ def extract_span_identifiers( def extract_tags( - opik_metadata: dict[str, Any], + opik_metadata: Mapping[str, Any], custom_llm_provider: str | None, ) -> list[str]: """ @@ -122,7 +123,7 @@ def apply_proxy_header_overrides( project_name: str, tags: list[str], thread_id: str | None, - proxy_headers: dict[str, Any], + proxy_headers: Mapping[str, str], ) -> tuple[str, list[str], str | None]: """ Apply overrides from proxy request headers (opik_* prefix). @@ -148,7 +149,7 @@ def apply_proxy_header_overrides( thread_id = value elif param_key == "tags": try: - parsed_tags = json.loads(value) + parsed_tags: object = json.loads(value) if isinstance(parsed_tags, list): tags.extend(parsed_tags) except (json.JSONDecodeError, TypeError): @@ -158,11 +159,11 @@ def apply_proxy_header_overrides( def extract_and_build_metadata( - opik_metadata: dict[str, Any], - standard_logging_metadata: dict[str, Any], - standard_logging_object: dict[str, Any], - litellm_kwargs: dict[str, Any], -) -> dict[str, Any]: + opik_metadata: Mapping[str, object], + standard_logging_metadata: Mapping[str, object], + standard_logging_object: Mapping[str, object], + litellm_kwargs: Mapping[str, object], +) -> dict[str, object]: """ Build the complete metadata dictionary from all available sources. diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d35405538f6..e8ed269f6cb 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -6,6 +6,7 @@ import json from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, cast from urllib.parse import urlsplit @@ -62,6 +63,31 @@ if TYPE_CHECKING: # --- typed sub-structures ---------------------------------------------------- # +def _cache_token_value(*values: object) -> int | None: + explicit_zero = False + invalid_before_zero = False + for raw_value in values: + if raw_value is None: + continue + if isinstance(raw_value, bool): + parsed = None + else: + try: + parsed = as_int(raw_value) + except (OverflowError, ValueError): + parsed = None + if parsed is None: + if not explicit_zero: + invalid_before_zero = True + elif parsed > 0: + return parsed + elif parsed == 0: + explicit_zero = True + elif not explicit_zero: + invalid_before_zero = True + return 0 if explicit_zero and not invalid_before_zero else None + + @dataclass(frozen=True) class LLMRequestParams: temperature: float | None = None @@ -104,12 +130,25 @@ class LLMUsage: metadata: Final[Mapping[str, object]] = payload.get("metadata") or {} raw_usage: Final = metadata.get("usage_object") usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {} + raw_details: Final = usage_object.get("prompt_tokens_details") + prompt_details: Final[Mapping[str, object]] = ( + raw_details if isinstance(raw_details, Mapping) else MappingProxyType({}) + ) return cls( input_tokens=as_int(payload.get("prompt_tokens")), output_tokens=as_int(payload.get("completion_tokens")), total_tokens=as_int(payload.get("total_tokens")), - cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")), - cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")), + cache_creation_input_tokens=_cache_token_value( + usage_object.get("cache_creation_input_tokens"), + prompt_details.get("cache_write_tokens"), + prompt_details.get("cache_creation_tokens"), + prompt_details.get("cache_creation_input_tokens"), + ), + cache_read_input_tokens=_cache_token_value( + usage_object.get("cache_read_input_tokens"), + prompt_details.get("cached_tokens"), + usage_object.get("prompt_cache_hit_tokens"), + ), ) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index c7e491c002a..e1623f4697f 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -11,9 +11,10 @@ identical metrics. The attribute cardinality filter is reused from v1 by import from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime -from typing import Any, Final, TypeAlias +from typing import Any, Final, Literal, Protocol, TypeAlias from opentelemetry.metrics import Histogram, Meter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -151,6 +152,29 @@ METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset( BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",) +class _TokenUsage(TypedDict, total=False): + """The token counts a response's ``usage`` carries, as the recorder reads them.""" + + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + + +class _ResponseView(Protocol): + """The one read the recorder makes on a litellm response object.""" + + def get(self, key: Literal["usage"], /) -> _TokenUsage | None: ... + + +class _MetricKwargs(TypedDict, total=False): + """The logging kwargs the recorder reads directly.""" + + call_type: ReadOnly[str | None] + litellm_params: ReadOnly[Mapping[str, object] | None] + response_cost: ReadOnly[float | None] + completion_start_time: ReadOnly[datetime | float | str | None] + api_call_start_time: ReadOnly[datetime | float | str | None] + + def resolve_error_type(kwargs: Mapping[str, Any]) -> str: """The ``error.type`` value for a failed request. @@ -192,8 +216,8 @@ class GenAIMetricRecorder: def record( self, - kwargs: Mapping[str, Any], - response_obj: Any, + kwargs: _MetricKwargs, + response_obj: _ResponseView | None, start_time: datetime, end_time: datetime, ) -> None: @@ -218,7 +242,7 @@ class GenAIMetricRecorder: def record_failure( self, - kwargs: Mapping[str, Any], + kwargs: _MetricKwargs, start_time: datetime, end_time: datetime, ) -> None: @@ -342,7 +366,7 @@ class GenAIMetricRecorder: # Per-metric recording # ------------------------------------------------------------------ # - def _record_token_usage(self, response_obj: Any, common_attrs: dict) -> None: + def _record_token_usage(self, response_obj: _ResponseView | None, common_attrs: dict) -> None: if not response_obj: return usage: Final = response_obj.get("usage") @@ -353,7 +377,7 @@ class GenAIMetricRecorder: self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs) self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs) - def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None: + def _record_time_to_first_token(self, kwargs: _MetricKwargs, common_attrs: dict) -> None: time_to_first_chunk: Final = time_to_first_chunk_seconds(kwargs) if time_to_first_chunk is None: return @@ -361,15 +385,14 @@ class GenAIMetricRecorder: def _record_time_per_output_token( self, - kwargs: Mapping[str, Any], - response_obj: Any, + kwargs: _MetricKwargs, + response_obj: _ResponseView | None, end_time: datetime, duration_s: float, common_attrs: dict, ) -> None: - completion_tokens = None - if response_obj and (usage := response_obj.get("usage")): - completion_tokens = usage.get("completion_tokens") + usage: Final = response_obj.get("usage") if response_obj else None + completion_tokens: Final = usage.get("completion_tokens") if usage else None if completion_tokens is None or completion_tokens <= 0: return diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index db9610a5a3c..4f7dff952e6 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -12,7 +12,10 @@ For batching specific details see CustomBatchLogger class import asyncio import atexit import os -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -34,6 +37,21 @@ from litellm.types.integrations.posthog import ( from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload +class PostHogBatchPayload(TypedDict): + api_key: ReadOnly[str] + batch: ReadOnly[Sequence[PostHogEventPayload]] + + +class PostHogLiteLLMParams(TypedDict, total=False): + metadata: ReadOnly[Mapping[str, object]] + + +class PostHogLogKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[StandardLoggingPayload] + standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams] + litellm_params: ReadOnly[PostHogLiteLLMParams] + + class PostHogLogger(CustomBatchLogger): def __init__(self, **kwargs): """ @@ -137,7 +155,7 @@ class PostHogLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() - def create_posthog_event_payload(self, kwargs: dict[str, Any]) -> PostHogEventPayload: + def create_posthog_event_payload(self, kwargs: PostHogLogKwargs) -> PostHogEventPayload: """ Helper function to create a PostHog event payload for logging @@ -171,11 +189,11 @@ class PostHogLogger(CustomBatchLogger): def _create_posthog_properties( self, standard_logging_object: StandardLoggingPayload, - kwargs: dict[str, Any], + kwargs: PostHogLogKwargs, event_name: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Create PostHog properties following LLM Analytics spec""" - properties: Final = {} + properties: Final[dict[str, object]] = {} # Core model information properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "") @@ -211,16 +229,19 @@ class PostHogLogger(CustomBatchLogger): properties["$ai_error"] = error_str # Add trace properties - self._add_trace_properties(properties, kwargs) + self._add_trace_properties(properties, standard_logging_object, kwargs) # Add custom metadata fields self._add_custom_metadata_properties(properties, kwargs) return properties - def _add_trace_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]): - standard_logging_object: Final = self._safe_get(kwargs, "standard_logging_object", {}) - + def _add_trace_properties( + self, + properties: dict[str, object], + standard_logging_object: StandardLoggingPayload, + kwargs: PostHogLogKwargs, + ) -> None: trace_id: Final = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid()) properties["$ai_trace_id"] = trace_id @@ -232,7 +253,7 @@ class PostHogLogger(CustomBatchLogger): if parent_id: properties["$ai_parent_id"] = parent_id - def _add_custom_metadata_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]): + def _add_custom_metadata_properties(self, properties: dict[str, object], kwargs: PostHogLogKwargs) -> None: """Add custom metadata fields to PostHog properties""" metadata: Final = self._extract_metadata(kwargs) if not isinstance(metadata, dict): @@ -277,7 +298,7 @@ class PostHogLogger(CustomBatchLogger): if key not in litellm_internal_fields: properties[key] = value - def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: dict[str, Any]) -> str: + def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: PostHogLogKwargs) -> str: metadata: Final = self._extract_metadata(kwargs) user_id: Final = self._safe_get(metadata, "user_id") if user_id: @@ -291,7 +312,7 @@ class PostHogLogger(CustomBatchLogger): return self._safe_uuid() - def _get_credentials_for_request(self, kwargs: dict[str, Any]) -> tuple[str | None, str | None]: + def _get_credentials_for_request(self, kwargs: PostHogLogKwargs) -> tuple[str | None, str | None]: """ Get PostHog credentials for this request. @@ -334,7 +355,7 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted") # Group events by credentials for batch sending - batches_by_credentials: Final[dict[tuple[str, str], list]] = {} + batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {} for item in self.log_queue: key = (item["api_key"], item["api_url"]) if key not in batches_by_credentials: @@ -380,18 +401,19 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.error("PostHog: Failed to initialize async components: %s", e) raise - def _extract_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]: - litellm_params: Final = kwargs.get("litellm_params", {}) or {} - return litellm_params.get("metadata", {}) or {} + def _extract_metadata(self, kwargs: PostHogLogKwargs) -> Mapping[str, object]: + litellm_params: Final[PostHogLiteLLMParams] = kwargs.get("litellm_params", {}) or {} + metadata: Final[Mapping[str, object]] = litellm_params.get("metadata", {}) or {} + return metadata def _safe_uuid(self) -> str: return str(uuid.uuid4()) - def _create_posthog_payload(self, events: list, api_key: str) -> dict[str, Any]: + def _create_posthog_payload(self, events: Sequence[PostHogEventPayload], api_key: str) -> PostHogBatchPayload: return {"api_key": api_key, "batch": events} - def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: - if obj is None or not hasattr(obj, "get"): + def _safe_get(self, obj: Mapping[str, object] | None, key: str, default: object = None) -> object: + if not isinstance(obj, Mapping): return default return obj.get(key, default) @@ -412,7 +434,7 @@ class PostHogLogger(CustomBatchLogger): try: # Group events by credentials (same logic as async_send_batch) - batches_by_credentials: Final[dict[tuple[str, str], list]] = {} + batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {} for item in self.log_queue: key = (item["api_key"], item["api_url"]) if key not in batches_by_credentials: diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 467ec72dc4a..975a9bd8639 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -8,6 +8,7 @@ import math import os import sys from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import replace from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast @@ -58,7 +59,10 @@ from litellm.types.utils import ( if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler + from prometheus_client import Gauge from prometheus_client.metrics import MetricWrapperBase + + from litellm.router import Router else: AsyncIOScheduler = Any @@ -67,6 +71,8 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel) _DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0 +UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other" + _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset( ( "guardrail_name", @@ -154,6 +160,44 @@ def _get_budget_metrics_per_request_timeout() -> float: return parsed +def _get_proxy_llm_router() -> Router | None: + try: + from litellm.proxy.proxy_server import llm_router + except Exception: + return None + return llm_router + + +def _bounded_requested_model_label(requested_model: str | None, router_originated: bool = False) -> str | None: + """ + Bound ``requested_model`` label cardinality: names the router recognizes + (model names, deployment ids, aliases, routing groups, team public model + names) or matches via a global or team wildcard/pattern route keep their + own label value; any other client-supplied string collapses into the + single ``other`` bucket. With no proxy router to vouch for the string, + client-supplied values collapse to ``other`` while ``router_originated`` + values (emitted by an SDK ``Router``'s own deployment failure and + fallback events, where the proxy router never exists) pass through. + """ + if not requested_model: + return requested_model + llm_router: Final = _get_proxy_llm_router() + if llm_router is None: + return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL + if llm_router.is_recognized_model(requested_model): + return requested_model + if requested_model in llm_router.team_public_model_names: + return requested_model + if llm_router.pattern_router.route(requested_model) is not None: + return requested_model + if any( + team_pattern_router.route(requested_model) is not None + for team_pattern_router in llm_router.team_pattern_routers.values() + ): + return requested_model + return UNRECOGNIZED_REQUESTED_MODEL_LABEL + + class PrometheusLogger(CustomLogger): # Class variables or attributes @@ -434,6 +478,30 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"), ) + self.litellm_api_key_rate_limit_allowed_metric = self._gauge_factory( + "litellm_api_key_rate_limit_allowed_metric", + "Configured rate limit for the API Key in the current window (rpm_limit / tpm_limit), by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_allowed_metric"), + ) + + self.litellm_api_key_rate_limit_used_metric = self._gauge_factory( + "litellm_api_key_rate_limit_used_metric", + "Requests or tokens the API Key has consumed in the current rate limit window, by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_used_metric"), + ) + + self.litellm_team_rate_limit_allowed_metric = self._gauge_factory( + "litellm_team_rate_limit_allowed_metric", + "Configured rate limit for the Team in the current window (team rpm_limit / tpm_limit), by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_team_rate_limit_allowed_metric"), + ) + + self.litellm_team_rate_limit_used_metric = self._gauge_factory( + "litellm_team_rate_limit_used_metric", + "Requests or tokens the Team has consumed in the current rate limit window, by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_team_rate_limit_used_metric"), + ) + ######################################## # LLM API Deployment Metrics / analytics ######################################## @@ -1433,6 +1501,11 @@ class PrometheusLogger(CustomLogger): model_id=enum_values.model_id, ) + self._set_key_and_team_rate_limit_metrics( + standard_logging_payload=standard_logging_payload, # pyright: ignore[reportArgumentType] # isinstance(dict) above narrows the TypedDict to dict[Unknown, Unknown] + enum_values=enum_values, + ) + # set latency metrics self._set_latency_metrics( kwargs=kwargs, @@ -1960,17 +2033,102 @@ class PrometheusLogger(CustomLogger): """ if standard_logging_payload is None: return None + return PrometheusLogger._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-model_per_key-remaining-{rate_limit_type}", + ) + + @staticmethod + def _get_int_from_v3_rate_limit_headers( + standard_logging_payload: StandardLoggingPayload, + header_name: str, + ) -> int | None: hidden_params: Final = standard_logging_payload.get("hidden_params") if hidden_params is None: return None - additional_headers: Final = hidden_params.get("additional_headers") + additional_headers: Final[Mapping[str, object] | None] = hidden_params.get("additional_headers") if additional_headers is None: return None - value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + value: Final = additional_headers.get(header_name) if isinstance(value, bool) or not isinstance(value, int): return None return value + def _set_key_and_team_rate_limit_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + ) -> None: + """ + Export the key-level and team-level RPM / TPM limit and current window + usage from the ``x-ratelimit-{api_key,team}-{limit,remaining}-*`` + headers the v3 rate limiter mirrors into the logging payload. The + limiter already read these counters (from Redis when configured) on + the request path, so no extra store lookup happens here. Descriptors + without a configured limit emit no header, so their series is removed + rather than left at the value from before the limit was dropped. + """ + descriptor_gauges: Final[ + tuple[tuple[Literal["api_key", "team"], DEFINED_PROMETHEUS_METRICS, Gauge, Gauge], ...] + ] = ( + ( + "api_key", + "litellm_api_key_rate_limit_allowed_metric", + self.litellm_api_key_rate_limit_allowed_metric, + self.litellm_api_key_rate_limit_used_metric, + ), + ( + "team", + "litellm_team_rate_limit_allowed_metric", + self.litellm_team_rate_limit_allowed_metric, + self.litellm_team_rate_limit_used_metric, + ), + ) + for descriptor_key, metric_name, allowed_gauge, used_gauge in descriptor_gauges: + for rate_limit_type in ("requests", "tokens"): + self._set_rate_limit_allowed_and_used_gauges( + standard_logging_payload=standard_logging_payload, + enum_values=enum_values, + descriptor_key=descriptor_key, + metric_name=metric_name, + allowed_gauge=allowed_gauge, + used_gauge=used_gauge, + rate_limit_type=rate_limit_type, + ) + + def _set_rate_limit_allowed_and_used_gauges( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + descriptor_key: Literal["api_key", "team"], + metric_name: DEFINED_PROMETHEUS_METRICS, + allowed_gauge: Gauge, + used_gauge: Gauge, + rate_limit_type: Literal["requests", "tokens"], + ) -> None: + limit: Final = self._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-{descriptor_key}-limit-{rate_limit_type}", + ) + remaining: Final = self._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-{descriptor_key}-remaining-{rate_limit_type}", + ) + labelled_values: Final = replace(enum_values, rate_limit_type=rate_limit_type) + labelnames: Final = self.get_labels_for_metric(metric_name) + labels: Final = prometheus_label_factory( + supported_enum_labels=labelnames, + enum_values=labelled_values, + label_context=PrometheusLabelFactoryContext(labelled_values), + ) + if limit is None or remaining is None: + label_values: Final = tuple(labels.get(label) for label in labelnames) + self._bounded_prometheus_series_tracker.remove_series(allowed_gauge, label_values) + self._bounded_prometheus_series_tracker.remove_series(used_gauge, label_values) + return + allowed_gauge.labels(**labels).set(limit) + used_gauge.labels(**labels).set(limit - remaining) + def _set_virtual_key_rate_limit_metrics( self, user_api_key: str | None, @@ -2407,7 +2565,7 @@ class PrometheusLogger(CustomLogger): team_alias=user_api_key_dict.team_alias, org_id=user_api_key_dict.org_id, org_alias=user_api_key_dict.organization_alias, - requested_model=request_data.get("model", ""), + requested_model=_bounded_requested_model_label(request_data.get("model", "")), status_code=str(status_code), exception_status=str(status_code), exception_class=self._get_exception_class_name(original_exception), @@ -2627,7 +2785,9 @@ class PrometheusLogger(CustomLogger): label_model_id = "" label_api_base = "" label_api_provider = "" - label_requested_model = litellm_model_name or model_group or "" + label_requested_model = ( + _bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or "" + ) enum_values: Final = UserAPIKeyLabelValues( litellm_model_name=label_litellm_model_name, @@ -3186,7 +3346,7 @@ class PrometheusLogger(CustomLogger): _tags: Final = cast(list[str], kwargs.get("tags") or []) enum_values: Final = UserAPIKeyLabelValues( - requested_model=original_model_group, + requested_model=_bounded_requested_model_label(original_model_group, router_originated=True), fallback_model=_new_model, hashed_api_key=standard_metadata["user_api_key_hash"], api_key_alias=standard_metadata["user_api_key_alias"], @@ -3227,7 +3387,7 @@ class PrometheusLogger(CustomLogger): ) enum_values: Final = UserAPIKeyLabelValues( - requested_model=original_model_group, + requested_model=_bounded_requested_model_label(original_model_group, router_originated=True), fallback_model=_new_model, hashed_api_key=standard_metadata["user_api_key_hash"], api_key_alias=standard_metadata["user_api_key_alias"], diff --git a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py index c54790b8ae7..c1ccf09d5d6 100644 --- a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py +++ b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py @@ -60,6 +60,10 @@ class BoundedPrometheusSeriesTracker: break del series[tracked_label_values] + def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool: + """Drop one child series, True when it is gone (removed or never existed).""" + return self._remove_metric_child(metric, label_values) + def _should_run_ttl_cleanup( self, metric_name: str, diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index ddeb410c54a..8ce461eea5b 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -1,11 +1,18 @@ #### What this does #### # On success + failure, log events to Supabase +import hashlib from datetime import datetime from typing import Final, cast import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.constants import ( + MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, + MAX_S3_OBJECT_KEY_BYTES, + S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, + S3_PREFIX_DIGEST_CHARS, +) from litellm.types.utils import StandardLoggingPayload @@ -133,9 +140,7 @@ class S3Logger: s3_file_name, ) - s3_object_download_filename: Final = ( - "time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json" - ) + s3_object_download_filename: Final = get_s3_object_download_filename(start_time, payload["id"]) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -198,6 +203,47 @@ def resolve_sse_params( return algorithm, valid_key_id +S3_MIN_BOUNDED_FILE_NAME_BYTES: Final = 64 + + +def _truncate_to_utf8_bytes(value: str, max_bytes: int) -> str: + """Trim `value` so its UTF-8 encoding fits `max_bytes`, never splitting a character.""" + if max_bytes <= 0: + return "" + encoded: Final = value.encode("utf-8") + if len(encoded) <= max_bytes: + return value + return encoded[:max_bytes].decode("utf-8", errors="ignore") + + +def get_s3_object_download_filename(start_time: datetime, response_id: str) -> str: + """Content-Disposition filename for the uploaded object, bounded to the metadata header cap.""" + sanitized_response_id: Final = response_id.replace("/", "_").replace('"', "_") + file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{response_id}" + sanitized_file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{sanitized_response_id}" + budget: Final = MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES - len(b".json") + if len(sanitized_file_name.encode("utf-8")) <= budget: + return sanitized_file_name + ".json" + return _bounded_s3_file_name(file_name, sanitized_file_name, budget) + ".json" + + +def _bounded_s3_file_name(s3_file_name: str, sanitized_s3_file_name: str, max_bytes: int) -> str: + """As much of the file name as `max_bytes` allows, then the sha256 of the whole name.""" + digest: Final = hashlib.sha256(s3_file_name.encode("utf-8")).hexdigest() + head_budget: Final = min(S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, max_bytes - len(digest) - 1) + head: Final = _truncate_to_utf8_bytes(sanitized_s3_file_name, head_budget) + return f"{head}_{digest}" if head else digest + + +def _bounded_s3_prefix(configured_prefix: str, max_bytes: int) -> str: + """As much of the configured prefix as fits, then a digest segment naming the full prefix.""" + digest_segment: Final = hashlib.sha256(configured_prefix.encode("utf-8")).hexdigest()[:S3_PREFIX_DIGEST_CHARS] + "/" + if max_bytes < len(digest_segment): + return "" + head: Final = _truncate_to_utf8_bytes(configured_prefix, max_bytes - len(digest_segment) - 1).rstrip("/") + return f"{head}/{digest_segment}" if head else digest_segment + + def get_s3_object_key( s3_path: str, prefix: str, @@ -205,12 +251,23 @@ def get_s3_object_key( s3_file_name: str, ) -> str: sanitized_s3_file_name: Final = s3_file_name.replace("/", "_") - s3_object_key = ( - (s3_path.rstrip("/") + "/" if s3_path else "") - + prefix - + start_time.strftime("%Y-%m-%d") - + "/" - + sanitized_s3_file_name - ) # we need the s3 key to include the time, so we log cache hits too - s3_object_key += ".json" - return s3_object_key + configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix + date_segment: Final = start_time.strftime("%Y-%m-%d") + "/" + # we need the s3 key to include the time, so we log cache hits too + s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json" + if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES: + return s3_object_key + + # shorten the response id first and only trim the configured prefix if that is what does not + # fit, so prefix scoped IAM policies and lifecycle rules keep matching + budget: Final = MAX_S3_OBJECT_KEY_BYTES - len(date_segment.encode("utf-8")) - len(b".json") + prefix_bytes: Final = len(configured_prefix.encode("utf-8")) + if prefix_bytes + S3_MIN_BOUNDED_FILE_NAME_BYTES <= budget: + bounded_file_name: Final = _bounded_s3_file_name(s3_file_name, sanitized_s3_file_name, budget - prefix_bytes) + return configured_prefix + date_segment + bounded_file_name + ".json" + + shortest_file_name: Final = _bounded_s3_file_name( + s3_file_name, sanitized_s3_file_name, S3_MIN_BOUNDED_FILE_NAME_BYTES + ) + bounded_prefix: Final = _bounded_s3_prefix(configured_prefix, budget - len(shortest_file_name.encode("utf-8"))) + return bounded_prefix + date_segment + shortest_file_name + ".json" diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 9f6ae72fb3a..712ce41d09e 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -16,7 +16,11 @@ from urllib.parse import quote import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS -from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params +from litellm.integrations.s3 import ( + get_s3_object_download_filename, + get_s3_object_key, + resolve_sse_params, +) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -259,11 +263,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): now: Final = datetime.now(timezone.utc) audit_log_id: Final = audit_log.get("id", "unknown") - s3_path = cast(str | None, self.s3_path) or "" - s3_path = s3_path.rstrip("/") + "/" if s3_path else "" - - s3_object_key: Final = ( - f"{s3_path}audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log_id}.json" + s3_object_key: Final = get_s3_object_key( + cast(str | None, self.s3_path) or "", + "audit_logs/", + now, + f"{now.strftime('%H-%M-%S')}_{audit_log_id}", ) element: Final = s3BatchLoggingElement( @@ -463,9 +467,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) verbose_logger.debug("s3_object_key=%s", s3_object_key) - s3_object_download_filename: Final = ( - f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" - ) + s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"]) return s3BatchLoggingElement( payload=dict(standard_logging_payload), diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 18bda0a9d55..27da785331a 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -1,8 +1,11 @@ """Shadow Eval Logger: samples a shadowed key's successful LLM requests (chat completions, Anthropic Messages, and Responses API surfaces, each normalized to chat shape), duplicates -each against the job's other arm in a detached task (the auto-router for a forward job, the -fixed baseline model for a reverse one), blind-judges real vs shadow, and appends one -``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write. +each through every shadow arm in one detached task (each candidate auto-router for a +forward job, the fixed baseline model for a reverse one), blind-judges real vs each arm, +and appends one ``LiteLLM_ShadowEvalAttempt`` row per arm (verdict or error) as the +feature's only hot-path write. A multi-router job's arms therefore score the identical +sampled requests against the identical real responses, which is what makes their win +rates comparable head-to-head. Counts, status, and spend derive from those rows at read time, so nothing can disagree across pods or stop races; the hook reads active jobs through a short-TTL cache.""" @@ -498,12 +501,16 @@ def _decision_classifier_cost(metadata: Mapping[str, object]) -> float: return float(raw) if isinstance(raw, (int, float)) else 0.0 -def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool: - """Whether the router under evaluation served this request, which is what decides - the direction it belongs to. A forward job skips its own router's traffic, since - duplicating it would compare the router to itself: guaranteed ties, judge spend for - zero information. A reverse job samples exactly that traffic and nothing else.""" - return _routing_decision(request_metadata).get("router_model_name") == router_name +def _direction_admits(request_metadata: Mapping[str, object], job: "ActiveShadowEvalJob") -> bool: + """Whether this request belongs to the job's direction. A forward job skips traffic + any of its candidate routers served: duplicating a router's own request compares it + to itself (guaranteed ties), and judging a sibling against another candidate's live + response would score candidates against each other instead of against the incumbent. + A reverse job samples exactly its one router's traffic and nothing else.""" + routed_by: Final = _routing_decision(request_metadata).get("router_model_name") + if job.direction == "reverse": + return routed_by == job.router_name + return routed_by not in job.arm_router_names @dataclass(frozen=True, slots=True) @@ -546,6 +553,7 @@ class ActiveShadowEvalJob(BaseModel): id: str router_name: str + router_names: tuple[str, ...] = () direction: ShadowEvalDirection = "forward" baseline_model: str | None = None shadow_percentage: float @@ -567,12 +575,25 @@ class ActiveShadowEvalJob(BaseModel): raise ValueError("baseline_model is set for exactly the reverse jobs") return self + @model_validator(mode="after") + def _reverse_evaluates_one_router(self) -> "ActiveShadowEvalJob": + """A reverse row naming several routers is unsamplable (there is no one traffic + slice they share) and fails closed.""" + if self.direction == "reverse" and len(self.arm_router_names) > 1: + raise ValueError("a reverse job evaluates exactly one router") + return self + @property - def shadow_target(self) -> str: - """The model the duplicated arm calls: the router itself for a forward job, the - fixed baseline for a reverse one. Total because the validator above pins + def arm_router_names(self) -> tuple[str, ...]: + """The job's full router set; rows from before router_names existed hold it in + router_name alone. The one place that reading lives on the sampling side.""" + return self.router_names or (self.router_name,) + + def arm_target(self, arm_router: str) -> str: + """The model one duplicated arm calls: the candidate router itself for a forward + job, the fixed baseline for a reverse one. Total because the validator above pins baseline_model to reverse jobs and only those.""" - return self.baseline_model or self.router_name + return self.baseline_model or arm_router def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowEvalJob | None: @@ -696,7 +717,7 @@ class ShadowEvalLogger(CustomLogger): now >= job.ends_at or job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns or (job.max_budget is not None and job.spend >= job.max_budget) - or _request_was_routed_by(request_metadata, job.router_name) != (job.direction == "reverse") + or not _direction_admits(request_metadata, job) ): continue if not _sample_hits(request_id, job.id, job.shadow_percentage): @@ -773,7 +794,10 @@ class ShadowEvalLogger(CustomLogger): if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS: self._record_funnel(job.id, "shed") continue - self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1 + # One start writes one attempt row per arm, and max_turns is a row + # ceiling, so admission must pre-count every arm or a multi-router + # job overshoots the valve N-fold within a cache generation. + self._job_starts[job.id] = self._job_starts.get(job.id, 0) + len(job.arm_router_names) self._inflight_shadow_tasks += 1 asyncio.create_task( self._run_shadow_eval( @@ -812,32 +836,74 @@ class ShadowEvalLogger(CustomLogger): shadow_params: Mapping[str, object], parent_metadata: Mapping[str, object], ) -> None: - """Budget gate -> shadow call -> blind judge -> one attempt row, and every exit - in exactly one coverage bucket: the gates that decline to spend on an admitted - sample (no DB to record into, an over-budget key, an unverifiable or exhausted - eval budget) count it withheld, so eligible traffic still reconciles as - not_sampled + unjudgeable + shed + withheld + attempt rows. The prisma gate sits - above the dispatch so no provider spend happens without a place to record the - outcome, and the budget read lives here rather than in the success hook.""" + """Budget gates once per sampled request, then every router arm in turn: shadow + call -> blind judge -> one attempt row stamped with the arm. The gates that + decline to spend on an admitted sample (no DB to record into, an over-budget key, + an unverifiable or exhausted eval budget) count the REQUEST withheld before any + arm runs, so funnel counters stay per-request and a leg's eligible traffic still + reconciles as not_sampled + unjudgeable + shed + withheld + sampled requests, + where each sampled request writes one attempt row per arm. A budget crossed + mid-loop lets the remaining arms overshoot by one round, the same class of + overshoot as the samples already in flight when the cap is crossed. The prisma + gate sits above the dispatch so no provider spend happens without a place to + record the outcome, and the budget read lives here rather than in the success + hook.""" prisma: Final = self._prisma_provider() + if prisma is None: + self._record_funnel(job.id, "withheld") + return + if await _key_or_team_is_over_budget(parent_metadata): + self._record_funnel(job.id, "withheld") + return + if job.max_budget is not None: + try: + spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget) + except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it + verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e) + self._record_funnel(job.id, "withheld") + return + if spend >= job.max_budget: + self._record_funnel(job.id, "withheld") + return + for arm_router in job.arm_router_names: + await self._run_shadow_arm( + prisma=prisma, + job=job, + arm_router=arm_router, + request_id=request_id, + messages=messages, + real_text=real_text, + real_model=real_model, + real_cost=real_cost, + real_classifier_cost=real_classifier_cost, + real_cache_hit=real_cache_hit, + control_tier=control_tier, + shadow_params=shadow_params, + parent_metadata=parent_metadata, + ) + + async def _run_shadow_arm( + self, + prisma: "PrismaClient", + job: ActiveShadowEvalJob, + arm_router: str, + request_id: str, + messages: Sequence[Mapping[str, object]], + real_text: str, + real_model: str, + real_cost: float, + real_classifier_cost: float, + real_cache_hit: bool, + control_tier: str | None, + shadow_params: Mapping[str, object], + parent_metadata: Mapping[str, object], + ) -> None: + """One arm's pipeline: shadow call -> blind judge -> one attempt row, every exit + recording this arm's outcome, so one arm's fault never silences a sibling arm.""" try: - if prisma is None: - self._record_funnel(job.id, "withheld") - return - if await _key_or_team_is_over_budget(parent_metadata): - self._record_funnel(job.id, "withheld") - return - if job.max_budget is not None: - try: - spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget) - except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it - verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e) - self._record_funnel(job.id, "withheld") - return - if spend >= job.max_budget: - self._record_funnel(job.id, "withheld") - return - shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata) + shadow: Final = await self._call_router_shadow( + job.arm_target(arm_router), messages, shadow_params, parent_metadata + ) except Exception as e: # noqa: BLE001 # detached task: nothing billed yet, record and never raise verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e) await self._record_attempt( @@ -845,6 +911,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=f"pipeline error: {e}", real_cost=real_cost, @@ -858,6 +925,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=shadow.error, shadow_cost=shadow.cost, @@ -882,6 +950,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=verdict.error, shadow=shadow, @@ -898,6 +967,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome=verdict.preference, shadow=shadow, real_model=real_model, @@ -916,6 +986,7 @@ class ShadowEvalLogger(CustomLogger): job, request_id, control_tier, + router_name=arm_router, outcome="error", error=f"pipeline error: {e}", shadow=shadow, @@ -933,6 +1004,7 @@ class ShadowEvalLogger(CustomLogger): request_id: str, control_tier: str | None, *, + router_name: str, outcome: str, real_cost: float, real_classifier_cost: float, @@ -955,6 +1027,7 @@ class ShadowEvalLogger(CustomLogger): data={ # mutable-ok: Prisma payload "job_id": job.id, "request_id": request_id, + "router_name": router_name, "outcome": outcome, "tier": control_tier if job.direction == "reverse" else (shadow.tier if shadow else None), "real_model": real_model or None, diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index aa29162ba1f..07d4f959489 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -13,7 +13,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.utils import StandardCallbackDynamicParams +from litellm.types.utils import CallTypes, StandardCallbackDynamicParams from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, VectorStoreResultContent, @@ -226,7 +226,7 @@ class VectorStorePreCallHook(CustomLogger): self, request_data: dict, response: Any, - call_type: Any | None, + call_type: CallTypes | None, ) -> Any | None: """ Add search results to the response after successful LLM call. @@ -283,7 +283,7 @@ class VectorStorePreCallHook(CustomLogger): self, request_data: dict, response_chunk: Any, - call_type: Any | None, + call_type: CallTypes | None, ) -> Any | None: """ Add search results to the final streaming chunk. diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 81310b9ddc3..dc61ee38a8c 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1633,16 +1633,18 @@ class WebSearchInterceptionLogger(CustomLogger): def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": if llm_router is None or not hasattr(llm_router, "search_tools"): return None - search_tools: Final = list(getattr(llm_router, "search_tools") or []) + search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ()) return self._select_search_tool_from_list(search_tools=search_tools, source="router") def _select_search_tool_from_list( self, - search_tools: list[_SearchToolConfig], + search_tools: Sequence[_SearchToolConfig], source: str, ) -> "_SearchToolConfig | None": if self.search_tool_name: - matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name] + matching_tools: Final = tuple( + tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name + ) if matching_tools: search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider") verbose_logger.debug( diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 3b3775a8fe6..dab3e48f91a 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -7,7 +7,13 @@ import os from dataclasses import dataclass from typing import Final -from litellm.types.files import get_file_mime_type_from_extension +from litellm.types.files import ( + AUDIO_FILE_TYPES, + FILE_EXTENSIONS, + FILE_MIME_TYPES, + FileType, + get_file_mime_type_from_extension, +) from litellm.types.utils import FileTypes @@ -323,3 +329,75 @@ def calculate_request_duration(file: FileTypes) -> float | None: except Exception: # Silently fail if duration extraction fails return None + + +DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg" + + +def _speech_media_type_for_response_format(response_format: str) -> str | None: + file_type: Final = next( + (candidate for candidate, extensions in FILE_EXTENSIONS.items() if response_format.lower() in extensions), + None, + ) + if file_type is None or file_type not in AUDIO_FILE_TYPES: + return None + return FILE_MIME_TYPES[file_type] + + +def resolve_speech_media_type(upstream_content_type: str | None, response_format: str | None) -> str: + upstream_media_type: Final = (upstream_content_type or "").split(";", 1)[0].strip().lower() + if upstream_media_type.startswith("audio/"): + return upstream_media_type + requested_media_type: Final = ( + None if response_format is None else _speech_media_type_for_response_format(response_format) + ) + return requested_media_type or DEFAULT_SPEECH_MEDIA_TYPE + + +_OGG_OPUS_HEAD_WINDOW: Final = 64 +_ADTS_SYNC_AND_LAYER_MASK: Final = 0xF6 +_ADTS_SYNC_AND_LAYER: Final = 0xF0 +_ADTS_SAMPLE_RATE_INDEX_LIMIT: Final = 13 +_MPEG_SYNC_MASK: Final = 0xE0 +_MPEG_LAYER_MASK: Final = 0x06 +_MPEG_RESERVED_VERSION: Final = 0x01 +_MPEG_INVALID_BITRATE_INDEX: Final = 0x0F +_MPEG_RESERVED_SAMPLE_RATE_INDEX: Final = 0x03 + + +def _adts_aac_frame_media_type(header: bytes) -> str | None: + sample_rate_index: Final = (header[2] >> 2) & 0x0F + return FILE_MIME_TYPES[FileType.AAC] if sample_rate_index < _ADTS_SAMPLE_RATE_INDEX_LIMIT else None + + +def _mpeg_audio_frame_media_type(header: bytes) -> str | None: + version: Final = (header[1] >> 3) & 0x03 + layer: Final = header[1] & _MPEG_LAYER_MASK + bitrate_index: Final = header[2] >> 4 + sample_rate_index: Final = (header[2] >> 2) & 0x03 + if ( + (header[1] & _MPEG_SYNC_MASK) != _MPEG_SYNC_MASK + or version == _MPEG_RESERVED_VERSION + or layer == 0 + or bitrate_index == _MPEG_INVALID_BITRATE_INDEX + or sample_rate_index == _MPEG_RESERVED_SAMPLE_RATE_INDEX + ): + return None + return FILE_MIME_TYPES[FileType.MP3] + + +def speech_media_type_from_audio_bytes(audio: bytes) -> str | None: + if audio[:4] == b"RIFF" and audio[8:12] == b"WAVE": + return FILE_MIME_TYPES[FileType.WAV] + if audio[:4] == b"fLaC": + return FILE_MIME_TYPES[FileType.FLAC] + if audio[:4] == b"OggS": + is_opus: Final = b"OpusHead" in audio[:_OGG_OPUS_HEAD_WINDOW] + return FILE_MIME_TYPES[FileType.OPUS if is_opus else FileType.OGG] + if audio[:3] == b"ID3": + return FILE_MIME_TYPES[FileType.MP3] + if len(audio) < 3 or audio[0] != 0xFF: + return None + if (audio[1] & _ADTS_SYNC_AND_LAYER_MASK) == _ADTS_SYNC_AND_LAYER: + return _adts_aac_frame_media_type(audio) + return _mpeg_audio_frame_media_type(audio) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 9b53b79bbe6..ce51fb19970 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider -def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None: +def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their @@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or model.split("/", 1)[0] + declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None @@ -536,6 +536,14 @@ def get_llm_provider( ) +def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig": + if custom_llm_provider == "qwencloud": + return litellm.QwenCloudChatConfig() + if custom_llm_provider == "qwen_ai_platform": + return litellm.QwenAIPlatformChatConfig() + return litellm.DashScopeChatConfig() + + def _get_openai_compatible_provider_info( model: str, api_base: str | None, @@ -785,11 +793,11 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, ) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info(api_base, api_key) - elif custom_llm_provider == "dashscope": + elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"): ( api_base, dynamic_api_key, - ) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) + ) = _dashscope_family_chat_config(custom_llm_provider)._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "modelscope": ( api_base, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 2043a9e2f89..9cba5db8ab7 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -12,6 +12,7 @@ import asyncio import json import os import random +import time from collections.abc import Awaitable, Callable from dataclasses import dataclass from datetime import datetime, timezone @@ -154,18 +155,6 @@ class GetModelCostMap: return True - @staticmethod - def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict: - """ - Fetch the model cost map from a remote URL. - - Returns the parsed JSON dict. Raises on network/parse errors - (caller is expected to handle). - """ - response: Final = httpx.get(url, timeout=timeout) - response.raise_for_status() - return response.json() - RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 @@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol): def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... +class _SyncGetClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ... + + +_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable + + def _default_reload_client() -> _AsyncGetClient: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient: return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap) -async def _attempt_fetch( - client: _AsyncGetClient, url: str, timeout: int -) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable: +def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome: + reason: Final = f"{type(error).__name__} fetching {url}: {error}" + if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)): + return ModelCostMapReloadUnavailable(reason=reason) + return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None) + + +async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: try: response: Final = await client.get(url, timeout=timeout) - except httpx.HTTPError as e: - return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None) + except (httpx.HTTPError, httpx.InvalidURL) as e: + return _classify_fetch_error(e, url) + return _classify_fetch_response(response, url) + + +def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: + try: + response: Final = client.get(url, timeout=timeout) + except (httpx.HTTPError, httpx.InvalidURL) as e: + return _classify_fetch_error(e, url) + return _classify_fetch_response(response, url) + + +def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome: if response.status_code in RETRYABLE_FETCH_STATUS_CODES: return _FetchAttemptRetryable( reason=f"HTTP {response.status_code} from {url}", @@ -242,6 +255,22 @@ async def _attempt_fetch( return ModelCostMapReloaded(model_cost_map=parsed) +def _next_retry_wait( + outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random +) -> float | ModelCostMapReloadUnavailable: + if attempt == max_attempts: + return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)") + wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng) + verbose_logger.warning( + "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", + attempt, + max_attempts, + outcome.reason, + wait_seconds, + ) + return wait_seconds + + async def _fetch_remote_model_cost_map_with_retry( url: str, timeout: int, @@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry( outcome = await _attempt_fetch(client=client, url=url, timeout=timeout) if not isinstance(outcome, _FetchAttemptRetryable): return outcome - if attempt == max_attempts: - return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)") - wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng) - verbose_logger.warning( - "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", - attempt, - max_attempts, - outcome.reason, - wait_seconds, - ) + wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) + if isinstance(wait_seconds, ModelCostMapReloadUnavailable): + return wait_seconds await sleep(wait_seconds) return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") +def _fetch_remote_model_cost_map_with_retry_sync( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, +) -> ModelCostMapReloadResult: + for attempt in range(1, max_attempts + 1): + outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) + if not isinstance(outcome, _FetchAttemptRetryable): + return outcome + wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) + if isinstance(wait_seconds, ModelCostMapReloadUnavailable): + return wait_seconds + sleep(wait_seconds) + return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") + + async def refetch_model_cost_map( url: str, timeout: int = 5, @@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) -def get_model_cost_map(url: str) -> dict: +def get_model_cost_map( + url: str, + timeout: int = 5, + max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, + sleep: Callable[[float], None] = time.sleep, + rng: random.Random | None = None, + client: "_SyncGetClient | None" = None, +) -> dict: """ Public entry point — returns the model cost map dict. 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. - 2. Otherwise fetches from ``url``, validates integrity, and falls back - to the local backup on any failure. + 2. Otherwise fetches from ``url``, retrying transient HTTP errors + (429/5xx/transport) with Retry-After-aware backoff, validates + integrity, and falls back to the local backup on any failure. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict: _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - try: - content: Final = GetModelCostMap.fetch_remote_model_cost_map(url) - except Exception as e: + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=rng if rng is not None else random.Random(), + client=client if client is not None else httpx, + ) + if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - str(e), + result.reason, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + content: Final = result.model_cost_map # Validate using cached count (cheap int comparison, no file I/O) if not GetModelCostMap.validate_model_cost_map( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 350e5403e4c..9a6fb11f978 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -111,6 +111,7 @@ from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.rerank import RerankResponse from litellm.types.utils import ( + DEPLOYMENT_SCOPED_PRICING_FIELDS, CachingDetails, CallTypes, CostBreakdown, @@ -255,6 +256,7 @@ _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggi # 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 sentry_sdk_instance = None capture_exception = None @@ -2957,13 +2959,25 @@ class Logging(LiteLLMLoggingBaseClass): "Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model ) self.model_call_details["response_cost"] = None + except Exception: # noqa: BLE001 # cost calculation must never block later callbacks (slot release) + verbose_logger.exception( + "Error calculating streaming response cost for model=%s. Setting 'response_cost' to None", + self.model, + ) + self.model_call_details["response_cost"] = None self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + try: + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) + except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release) + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload " + "for a streaming response; callbacks still run without it" + ) # print standard logging payload if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: @@ -3003,32 +3017,39 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## for callback in callbacks: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks + try: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, ) - is not True - ): - continue - - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, + elif isinstance(callback, CustomLogger): + result = redact_message_input_output_from_custom_logger( + result=result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", + traceback.format_exc(), ) + self._handle_callback_failure(callback=callback) self.has_run_logging(event_type="async_success") @@ -5033,7 +5054,9 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool: """ Check if the model uses custom pricing - Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info` + Returns True if any custom pricing field is present in `litellm_params`, or if + any custom pricing or deployment-scoped pricing field (such as + ``off_peak_pricing``) is present in the metadata ``model_info`` """ if litellm_params is None: return False @@ -5051,7 +5074,7 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool: model_info: dict = metadata.get("model_info", {}) or {} if model_info: - matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys() + matching_keys = _MODEL_INFO_CUSTOM_PRICING_KEYS & model_info.keys() for key in matching_keys: if model_info.get(key) is not None: return True diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 9250b92e268..5504756ceb8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -3,7 +3,7 @@ Helper utilities for tracking the cost of built-in tools. """ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS @@ -16,6 +16,7 @@ from litellm.types.llms.openai import ( WebSearchOptions, ) from litellm.types.utils import ( + ChatCompletionAnnotation, Message, ModelInfo, ModelResponse, @@ -49,7 +50,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def get_cost_for_built_in_tools( model: str, - response_object: Any, + response_object: object, usage: Usage | None = None, custom_llm_provider: str | None = None, standard_built_in_tools_params: StandardBuiltInToolsParams | None = None, @@ -201,8 +202,7 @@ class StandardBuiltInToolCostTracking: model_info: Final = StandardBuiltInToolCostTracking._safe_get_model_info( model=model, custom_llm_provider=custom_llm_provider ) - file_search_raw: Final[Any] = standard_built_in_tools_params.get("file_search", {}) - file_search_usage: Final[FileSearchTool | None] = FileSearchTool(**file_search_raw) if file_search_raw else None + file_search_usage: Final[FileSearchTool | None] = standard_built_in_tools_params.get("file_search") or None # Convert model_info to dict and extract usage parameters model_info_dict: Final = dict(model_info) if model_info is not None else None @@ -245,7 +245,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def _extract_file_search_params( - file_search_usage: Any, + file_search_usage: object, ) -> tuple[float | None, float | None]: """Extract and convert file search parameters safely.""" storage_gb = None @@ -335,7 +335,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def _extract_token_counts( - computer_use_usage: Any, + computer_use_usage: object, ) -> tuple[int | None, int | None]: """Extract and convert token counts safely.""" input_tokens = None @@ -351,9 +351,9 @@ class StandardBuiltInToolCostTracking: return input_tokens, output_tokens @staticmethod - def _safe_convert_to_int(value: Any) -> int | None: + def _safe_convert_to_int(value: object) -> int | None: """Safely convert a value to int.""" - if value is not None: + if isinstance(value, (int, float, str)): try: return int(value) except (TypeError, ValueError): @@ -381,7 +381,7 @@ class StandardBuiltInToolCostTracking: return usage.model_copy(update={"server_tool_use": server_tool_use}) @staticmethod - def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool: + def response_object_includes_web_search_call(response_object: object, usage: Usage | None = None) -> bool: """ Check if the response object includes a web search call. @@ -446,7 +446,7 @@ class StandardBuiltInToolCostTracking: @staticmethod def response_object_includes_file_search_call( - response_object: Any, + response_object: object, ) -> bool: """ Check if the response object includes a file search call. @@ -477,11 +477,11 @@ class StandardBuiltInToolCostTracking: message: Message | None = getattr(choice, "message", None) if message is None: continue - if annotations := getattr(message, "annotations", None): - if len(annotations) > 0: - for annotation in annotations: - if annotation.get("type", None) == annotation_type: - return True + annotations: list[ChatCompletionAnnotation] | None = getattr(message, "annotations", None) + if annotations: + for annotation in annotations: + if annotation.get("type", None) == annotation_type: + return True return False @staticmethod @@ -522,10 +522,8 @@ class StandardBuiltInToolCostTracking: if model_info is None: return 0.0 - search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {}) - search_context_pricing: Final[SearchContextCostPerQuery] = ( - SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() - ) + search_context_raw: Final = model_info.get("search_context_cost_per_query") + search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery() if web_search_options.get("search_context_size", None) == "low": return search_context_pricing.get("search_context_size_low", 0.0) elif web_search_options.get("search_context_size", None) == "medium": @@ -545,10 +543,8 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 - search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {}) or {} - search_context_pricing: Final[SearchContextCostPerQuery] = ( - SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery() - ) + search_context_raw: Final = model_info.get("search_context_cost_per_query") + search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery() return search_context_pricing.get("search_context_size_medium", 0.0) @staticmethod @@ -714,7 +710,7 @@ class StandardBuiltInToolCostTracking: response_object: ModelResponse, ) -> bool: for _choice in response_object.choices: - message = getattr(_choice, "message", None) + message: Message | None = getattr(_choice, "message", None) if ( message is not None and hasattr(message, "annotations") diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 19e3f624268..b34c416cd40 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -2,10 +2,12 @@ ## Helper utilities for cost_per_token() import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timezone, tzinfo from types import MappingProxyType from typing import Any, Final, Literal, TypedDict, cast +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import litellm from litellm._logging import verbose_logger @@ -290,10 +292,187 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float, ) +def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool: + """Return True if current_time (UTC, defaulting to now) falls inside any off-peak window. + + off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers + with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past + midnight, and a window whose start equals its end covers the whole day. The start is + inclusive and the end is exclusive; malformed windows are ignored. + + An aware current_time is converted to UTC. A naive one is taken to already be UTC rather + than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(), + or every window shifts by the host's offset. + """ + reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time() + windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc + for window in windows: + try: + start_str, end_str = window.split("-") + start = datetime.strptime(start_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time() + end = datetime.strptime(end_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time() + except (ValueError, AttributeError): + continue + if start < end: + if start <= now < end: + return True + elif now >= start or now < end: + return True + return False + + +_WEEKDAY_NUMBERS: Final = MappingProxyType( + { + "mon": 1, + "monday": 1, + "tue": 2, + "tues": 2, + "tuesday": 2, + "wed": 3, + "wednesday": 3, + "thu": 4, + "thur": 4, + "thurs": 4, + "thursday": 4, + "fri": 5, + "friday": 5, + "sat": 6, + "saturday": 6, + "sun": 7, + "sunday": 7, + } +) + + +def _normalize_weekday(value: object) -> int | None: + if isinstance(value, bool): + return None + if isinstance(value, int): + return value if 1 <= value <= 7 else None + if isinstance(value, str): + return _WEEKDAY_NUMBERS.get(value.strip().lower()) + return None + + +def _weekday_calendar(weekday_timezone: object) -> tzinfo: + if isinstance(weekday_timezone, str) and weekday_timezone.strip(): + try: + return ZoneInfo(weekday_timezone.strip()) + except (ValueError, ZoneInfoNotFoundError): + return timezone.utc + return timezone.utc + + +def _matches_weekdays(reference_utc: datetime, weekdays: object, weekday_timezone: object) -> bool: + """Return True when reference_utc falls on one of the rule's weekdays, read on the calendar + named by weekday_timezone (default UTC). An absent weekdays means every day. The calendar + matters even when UTC and vendor-local weekdays agree at every currently priced hour: a + window past 16:00 UTC is where an Asia/Shanghai weekday diverges from the UTC one. + """ + if weekdays is None: + return True + if isinstance(weekdays, str) or not isinstance(weekdays, Sequence): + return False + allowed: Final = frozenset(day for day in map(_normalize_weekday, weekdays) if day is not None) + return reference_utc.astimezone(_weekday_calendar(weekday_timezone)).isoweekday() in allowed + + +def _as_window_strings(value: object) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, Sequence): + return tuple(entry for entry in value if isinstance(entry, str)) + return () + + +def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = None) -> bool: + """Return True when current_time (UTC, defaulting to now) is off-peak under the block's + rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose + hours apply only on its weekdays. + """ + reference: Final = current_time if current_time is not None else datetime.now(timezone.utc) + reference_utc: Final = ( + reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc) + ) + flat_windows: Final = _as_window_strings(off_peak.get("hours_utc")) + if flat_windows and _is_within_off_peak_window(flat_windows, reference_utc): + return True + windows: Final = off_peak.get("windows") + if isinstance(windows, str) or not isinstance(windows, Sequence): + return False + weekday_timezone: Final = off_peak.get("weekday_timezone") + for rule in windows: + if not isinstance(rule, Mapping): + continue + rule_windows = _as_window_strings(rule.get("hours_utc")) + if not rule_windows: + continue + if not _matches_weekdays(reference_utc, rule.get("weekdays"), weekday_timezone): + continue + if _is_within_off_peak_window(rule_windows, reference_utc): + return True + return False + + +def _coerce_off_peak_rate(value: object, default: float) -> float: + if isinstance(value, bool): + return default + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + return default + return default + + +def _apply_off_peak_pricing( + model_info: ModelInfo, + current_time: datetime | None, + prompt_base_cost: float, + completion_base_cost: float, + cache_read_cost: float, +) -> tuple[float, float, float]: + """Swap in off-peak per-token rates when the current UTC time is inside one of the model's + off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in + windows. An off-peak rate replaces the rate that would otherwise apply rather than + discounting it, so a model that also has tiered or above-threshold pricing bills the flat + off-peak rate for the whole request while the window is open. Any rate left unset in + off_peak_pricing falls back to the standard rate. + """ + off_peak: Final = model_info.get("off_peak_pricing") + if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time): + return prompt_base_cost, completion_base_cost, cache_read_cost + return ( + _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost), + _coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost), + _coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost), + ) + + +def _apply_off_peak_to_base_costs( + model_info: ModelInfo, + current_time: datetime | None, + base_costs: tuple[float, float, float, float, float], +) -> tuple[float, float, float, float, float]: + """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path + produced them. Cache-creation rates are passed through untouched, since off_peak_pricing + has no field for them. + """ + prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs + off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing( + model_info, current_time, prompt, completion, cache_read + ) + return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) + + def _get_token_base_cost( model_info: ModelInfo, usage: Usage, service_tier: str | None = None, + current_time: datetime | None = None, *, threshold_is_inclusive: bool = False, ) -> tuple[float, float, float, float, float]: @@ -311,7 +490,7 @@ def _get_token_base_cost( """ tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage) if tiered_base_costs is not None: - return tiered_base_costs + return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs) # Get service tier aware cost keys input_cost_key: Final = _get_service_tier_cost_key("input_cost_per_token", service_tier) @@ -345,12 +524,16 @@ def _get_token_base_cost( k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES) ] if not threshold_keys: - return ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, + return _apply_off_peak_to_base_costs( + model_info, + current_time, + ( + prompt_base_cost, + completion_base_cost, + cache_creation_cost, + cache_creation_cost_above_1hr, + cache_read_cost, + ), ) # Only sort the threshold keys (typically 1-2 keys instead of 66+) @@ -451,12 +634,16 @@ def _get_token_base_cost( except Exception: continue - return ( - prompt_base_cost, - completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, - cache_read_cost, + return _apply_off_peak_to_base_costs( + model_info, + current_time, + ( + prompt_base_cost, + completion_base_cost, + cache_creation_cost, + cache_creation_cost_above_1hr, + cache_read_cost, + ), ) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index bda365102af..ff46440ff5c 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -555,10 +555,10 @@ def update_messages_with_model_file_ids( def update_responses_input_with_model_file_ids( - input: Any, + input: object, model_id: str | None = None, model_file_id_mapping: dict[str, dict[str, str]] | None = None, -) -> str | list[dict[str, Any]]: +) -> object: """ Updates responses API input with provider-specific file IDs. File IDs are always inside the content array, not as direct input_file items. @@ -639,8 +639,8 @@ def update_responses_input_with_model_file_ids( def _decode_vector_store_ids_in_tools( - tools: list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: + tools: list[dict[str, object]] | None, +) -> list[dict[str, object]] | None: """ Decodes unified (LiteLLM-managed) vector_store_ids in file_search tools to provider-native IDs. Non-unified IDs are passed through unchanged. @@ -692,10 +692,10 @@ def _decode_vector_store_ids_in_tools( def update_responses_tools_with_model_file_ids( - tools: list[dict[str, Any]] | None, + tools: list[dict[str, object]] | None, model_id: str | None = None, model_file_id_mapping: dict[str, dict[str, str]] | None = None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """ Updates responses API tools with provider-specific file IDs. @@ -888,7 +888,7 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData: # --------------------------------------------------------------------------- -def _estimate_json_bytes(obj: Any) -> int: +def _estimate_json_bytes(obj: object) -> int: """Estimate the JSON-serialised byte size of ``obj`` without materialising JSON. Walks iteratively (no recursion stack risk). @@ -1281,6 +1281,19 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo +def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = tool.get("function") + if not isinstance(function, dict): + return tool + parameters: Final = function.get("parameters") + if not isinstance(parameters, dict): + return tool + flattened: Final = flatten_top_level_schema_combinators(parameters) + if flattened is parameters: + return tool + return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + + def _get_image_mime_type_from_url(url: str) -> str | None: """ Get mime type for common image URLs @@ -1979,7 +1992,7 @@ def drop_tool_reference_parts_from_tool_messages( return [_drop_tool_reference_parts(message) for message in messages] # mutable-ok: pipelines mutate message lists -def _attempt_json_repair(s: str) -> Any | None: +def _attempt_json_repair(s: str) -> object | None: """ Attempt to repair truncated JSON produced by LLM tool calls. @@ -2095,7 +2108,7 @@ def parse_tool_call_arguments( raise ValueError(error_message) from original_error -def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: +def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: """ Split a string that contains one or more concatenated JSON objects into a list of parsed dicts. @@ -2131,7 +2144,7 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: return [] decoder: Final = json.JSONDecoder() - results: Final[list[dict[str, Any]]] = [] + results: Final[list[dict[str, object]]] = [] idx = 0 length: Final = len(raw) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 795fb36961e..ba59e3fa997 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1694,6 +1694,18 @@ def convert_function_to_anthropic_tool_invoke( raise e +def _find_server_tool_result( + tool_id: str, + web_search_results: Sequence[object] | None, + tool_results: Sequence[object] | None, +) -> dict[str, object] | None: + candidates: Final = (*(web_search_results or ()), *(tool_results or ())) + return next( + (result for result in candidates if isinstance(result, dict) and result.get("tool_use_id") == tool_id), + None, + ) + + def convert_to_anthropic_tool_invoke( tool_calls: list[ChatCompletionAssistantToolCall], web_search_results: list[Any] | None = None, @@ -1758,32 +1770,22 @@ def convert_to_anthropic_tool_invoke( context="Anthropic tool invoke", ) - # Check if this is a server-side tool (web_search, tool_search, etc.) - # Server tool IDs start with "srvtoolu_" - if tool_id.startswith("srvtoolu_"): - # Create server_tool_use block instead of tool_use - _anthropic_server_tool_use: dict[str, object] = { - "type": "server_tool_use", - "id": tool_id, - "name": tool_name, - "input": tool_input, - } - anthropic_tool_invoke.append(_anthropic_server_tool_use) - - # Add corresponding tool result if available. - # Check both web_search_results (web_search_tool_result / web_fetch_tool_result) - # and tool_results (bash_code_execution_tool_result, etc.) - _all_tool_results: list[Any] = [] - if web_search_results: - _all_tool_results.extend(web_search_results) - if tool_results: - _all_tool_results.extend(tool_results) - for result in _all_tool_results: - if result.get("tool_use_id") == tool_id: - anthropic_tool_invoke.append(result) - break + server_tool_result = ( + _find_server_tool_result(tool_id, web_search_results, tool_results) + if tool_id.startswith("srvtoolu_") + else None + ) + if server_tool_result is not None: + anthropic_tool_invoke.append( + { + "type": "server_tool_use", + "id": tool_id, + "name": tool_name, + "input": tool_input, + } + ) + anthropic_tool_invoke.append(server_tool_result) else: - # Regular tool_use sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id) _anthropic_tool_use_param = AnthropicMessagesToolUseParam( type="tool_use", @@ -4955,10 +4957,13 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None: - from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock + from litellm.llms.bedrock.common_utils import ( + bedrock_model_accepts_cache_points, + is_claude_4_5_on_bedrock, + ) cache_control: Final = tool.get("cache_control", None) - if cache_control is not None: + if cache_control is not None and bedrock_model_accepts_cache_points(model): cache_point: Final = cache_control.get("type", "ephemeral") if cache_point == "ephemeral": cache_point_block: Final[CachePointBlock] = {"type": "default"} diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 9125ed6e70a..8479e108d17 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1500,6 +1500,6 @@ class RealTimeStreaming: pass -def client_sent_openai_beta_realtime_header(websocket: Any) -> bool: +def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: """True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``.""" return RealTimeStreaming._detect_beta_header(websocket) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e2139d688b..0e01577b20e 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -36,6 +36,8 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from openai.types.completion_usage import CompletionUsage + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -73,6 +75,18 @@ class _ContentChunk(TypedDict): choices: Sequence[_ContentChoice] +class _FunctionCallDelta(TypedDict): + function_call: ReadOnly[FunctionCall] + + +class _FunctionCallChoice(TypedDict): + delta: ReadOnly[_FunctionCallDelta] + + +class _FunctionCallChunk(TypedDict): + choices: ReadOnly[Sequence[_FunctionCallChoice]] + + class _AudioDelta(TypedDict, total=False): audio: ChatCompletionAudioDelta | None @@ -588,7 +602,7 @@ class ChunkProcessor: return tool_calls_list - def get_combined_function_call_content(self, function_call_chunks: list[dict[str, Any]]) -> FunctionCall: + def get_combined_function_call_content(self, function_call_chunks: Sequence["_FunctionCallChunk"]) -> FunctionCall: argument_list: Final = [] delta = function_call_chunks[0]["choices"][0]["delta"] function_call = delta.get("function_call", "") @@ -782,7 +796,7 @@ class ChunkProcessor: @staticmethod def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: - usage_chunk: Usage | None = None + usage_chunk: Usage | CompletionUsage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -794,7 +808,9 @@ class ChunkProcessor: if isinstance(usage_chunk, dict): return Usage(**usage_chunk) - return usage_chunk + if usage_chunk is None or isinstance(usage_chunk, Usage): + return usage_chunk + return Usage(**usage_chunk.model_dump()) def _calculate_usage_per_chunk( self, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f5f671e7585..480b1921c18 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -862,6 +862,8 @@ class CustomStreamWrapper: model_response: Final = ModelResponseStream(**args) if self.response_id is not None: model_response.id = self.response_id + elif model_response.id: + self.response_id = model_response.id if self.system_fingerprint is not None: model_response.system_fingerprint = self.system_fingerprint diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index c350bc5569e..3732ffd734c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -4,8 +4,9 @@ import base64 import io import struct from collections.abc import Callable, Iterable, Mapping, Sequence -from typing import Any, Final, Literal, cast +from typing import Final, Literal, cast +import httpx import tiktoken import litellm @@ -171,6 +172,10 @@ def calculate_tiles_needed( return total_tiles +def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: + return struct.unpack(fmt, buffer) + + def get_image_type(image_data: bytes) -> str | None: """take an image (really only the first ~100 bytes max are needed) and return 'png' 'gif' 'jpeg' 'webp' 'heic' or None. method added to @@ -210,9 +215,9 @@ def get_image_dimensions( if data.startswith(("http://", "https://")): try: client: Final = _get_httpx_client() - response: Final = safe_get(client, data) + response: Final[httpx.Response] = safe_get(client, data) max_bytes: Final = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) - content_length: Final = response.headers.get("Content-Length") + content_length: Final[str | None] = response.headers.get("Content-Length") if content_length is not None and int(content_length) > max_bytes: pass # skip download; img_data stays None else: @@ -229,10 +234,10 @@ def get_image_dimensions( img_type: Final = get_image_type(img_data) if img_type == "png": - w, h = struct.unpack(">LL", img_data[16:24]) + w, h = _unpack_ints(">LL", img_data[16:24]) return w, h elif img_type == "gif": - w, h = struct.unpack("H", fhandle.read(2))[0] - 2 + size = _unpack_ints(">H", fhandle.read(2))[0] - 2 fhandle.seek(1, 1) - h, w = struct.unpack(">HH", fhandle.read(4)) + h, w = _unpack_ints(">HH", fhandle.read(4)) return w, h elif img_type == "webp": # For WebP, the dimensions are stored at different offsets depending on the format # Check for VP8X (extended format) if img_data[12:16] == b"VP8X": - w = struct.unpack("> 14) & 0x3FFF) + 1 return w, h @@ -420,8 +425,8 @@ def token_counter( def _count_function_call_tokens( key: str, - value: Any, - message: Mapping[str, Any], + value: object, + message: Mapping[str, object], count_function: TokenCounterFunction, ) -> int: """ @@ -587,7 +592,7 @@ def _fix_model_name(model: str) -> str: def _count_image_tokens( - image_url: Any, + image_url: object, use_default_image_token_count: bool, ) -> int: """ @@ -627,7 +632,7 @@ def _count_image_tokens( raise ValueError(f"Invalid image_url type: {type(image_url).__name__}. Expected str or dict with 'url' field.") -def _validate_anthropic_content(content: Mapping[str, Any]) -> type: +def _validate_anthropic_content(content: Mapping[str, object]) -> type: """ Validate and determine which Anthropic TypedDict applies. @@ -642,7 +647,7 @@ def _validate_anthropic_content(content: Mapping[str, Any]) -> type: "tool_result": AnthropicMessagesToolResultParam, } - expected_cls: Final = mapping.get(content_type) + expected_cls: Final = mapping.get(content_type) if isinstance(content_type, str) else None if expected_cls is None: raise ValueError(f"Unknown Anthropic content type: '{content_type}'") @@ -714,7 +719,7 @@ def _count_file_tokens( def _count_anthropic_content( - content: Mapping[str, Any], + content: Mapping[str, object], count_function: TokenCounterFunction, use_default_image_token_count: bool, default_token_count: int | None, @@ -729,7 +734,7 @@ def _count_anthropic_content( avoiding hardcoded field names. """ typeddict_cls: Final = _validate_anthropic_content(content) - type_hints: Final = getattr(typeddict_cls, "__annotations__", {}) + type_hints: Final[Mapping[str, object]] = getattr(typeddict_cls, "__annotations__", {}) tokens = 0 # Fields to skip (metadata/identifiers that don't contribute to prompt tokens) diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 1c5ba951942..f1c7451796d 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -11,8 +11,11 @@ A2A Protocol Format: """ import json +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Optional +from typing_extensions import ReadOnly, TypedDict + from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.types.utils import GenericGuardrailAPIInputs @@ -23,6 +26,13 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth +class _A2ATextPart(TypedDict, total=False): + """The subset of an A2A message part this handler reads text from.""" + + kind: ReadOnly[str] + text: ReadOnly[str] + + class A2AGuardrailHandler(BaseTranslation): """ Handler for processing A2A Protocol messages with guardrails. @@ -41,7 +51,7 @@ class A2AGuardrailHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, - ) -> Any: + ) -> dict: """ Process A2A input messages by applying guardrails to text content. @@ -214,12 +224,12 @@ class A2AGuardrailHandler(BaseTranslation): async def process_output_streaming_response( self, - responses_so_far: list[Any], + responses_so_far: list[object], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, - ) -> list[Any]: + ) -> list[object]: """ Process A2A streaming output by applying guardrails to accumulated text. @@ -305,11 +315,12 @@ class A2AGuardrailHandler(BaseTranslation): def _parse_streaming_responses( self, - responses_so_far: list[Any], - ) -> tuple[list[dict[str, Any] | None], list[tuple[int, dict[str, Any]]]]: + responses_so_far: list[object], + ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" - parsed: Final[list[dict[str, Any] | None]] = [None] * len(responses_so_far) + parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) for i, item in enumerate(responses_so_far): + obj: dict[str, object] if isinstance(item, dict): obj = item elif isinstance(item, str): @@ -326,7 +337,7 @@ class A2AGuardrailHandler(BaseTranslation): def _collect_text_from_parsed_chunks( self, - valid_parsed: list[tuple[int, dict[str, Any]]], + valid_parsed: list[tuple[int, dict[str, object]]], ) -> tuple[str, list[int]]: """Collect text from parsed chunks, returning combined text and indices.""" from litellm.llms.a2a.common_utils import extract_text_from_a2a_response @@ -411,7 +422,7 @@ class A2AGuardrailHandler(BaseTranslation): def _extract_texts_from_parts( self, - parts: list[dict[str, Any]], + parts: Sequence[_A2ATextPart], path: tuple[str, ...], texts_to_check: list[str], task_mappings: list[tuple[tuple[str, ...], int]], diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9395cc3d33e..c7d12e5cf3a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -100,16 +100,6 @@ InputWriteBackTarget = ( ) -class _SSEDelta(TypedDict, total=False): - type: ReadOnly[str] - text: ReadOnly[str] - stop_reason: ReadOnly[str | None] - - -class _SSEEventData(TypedDict, total=False): - delta: ReadOnly[_SSEDelta] - - def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: return value @@ -157,6 +147,16 @@ class ExtractedInput: EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) +class _AnthropicSSEDelta(TypedDict, total=False): + type: ReadOnly[str] + text: ReadOnly[str] + stop_reason: ReadOnly[str | None] + + +class _AnthropicSSEEvent(TypedDict, total=False): + delta: ReadOnly[_AnthropicSSEDelta] + + class AnthropicMessagesHandler(BaseTranslation): """Process Anthropic messages with guardrails. @@ -1247,8 +1247,8 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data: _SSEEventData = json.loads(data_line) - delta = data.get("delta", {}) + data: _AnthropicSSEEvent = json.loads(data_line) + delta: _AnthropicSSEDelta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") except json.JSONDecodeError: @@ -1310,9 +1310,9 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data: _SSEEventData = json.loads(data_line) - delta = data.get("delta", {}) - stop_reason = delta.get("stop_reason") + data: _AnthropicSSEEvent = json.loads(data_line) + delta: _AnthropicSSEDelta = data.get("delta", {}) + stop_reason: str | None = delta.get("stop_reason") if stop_reason is not None: return True except json.JSONDecodeError: diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index cd47cdd57d6..c82be07a5c5 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -66,6 +66,10 @@ if TYPE_CHECKING: from litellm.llms.base_llm.chat.transformation import BaseConfig +def _loads_stream_chunk(payload: str) -> dict[str, object]: + return json.loads(payload) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -78,7 +82,7 @@ async def make_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, -) -> tuple[Any, httpx.Headers]: +) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_aclient @@ -93,7 +97,7 @@ async def make_call( ) except httpx.HTTPStatusError as e: error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise AnthropicError( @@ -138,7 +142,7 @@ def make_sync_call( json_mode: bool, speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, -) -> tuple[Any, httpx.Headers]: +) -> tuple["ModelResponseIterator", httpx.Headers]: if client is None: client = litellm.module_level_client # re-use a module level client @@ -153,7 +157,7 @@ def make_sync_call( ) except httpx.HTTPStatusError as e: error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) raise AnthropicError( @@ -292,7 +296,7 @@ class AnthropicChatCompletion(BaseLLM): status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) if error_response and hasattr(error_response, "text"): @@ -593,7 +597,7 @@ class AnthropicChatCompletion(BaseLLM): status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) error_text = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) if error_response and hasattr(error_response, "text"): @@ -664,10 +668,10 @@ class ModelResponseIterator: # Accumulate web_search_tool_result blocks for multi-turn reconstruction # See: https://github.com/BerriAI/litellm/issues/17737 - self.web_search_results: list[dict[str, Any]] = [] + self.web_search_results: list[dict[str, object]] = [] # Accumulate compaction blocks for multi-turn reconstruction - self.compaction_blocks: list[dict[str, Any]] = [] + self.compaction_blocks: list[dict[str, object]] = [] # Accumulate streamed thinking text so final usage can split reasoning # tokens from regular output tokens. @@ -727,7 +731,7 @@ class ModelResponseIterator: str, ChatCompletionToolCallChunk | None, list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], - dict[str, Any], + dict[str, object], str | None, ]: """ @@ -735,7 +739,7 @@ class ModelResponseIterator: """ text = "" tool_use: ChatCompletionToolCallChunk | None = None - provider_specific_fields: Final = {} + provider_specific_fields: Final[dict[str, object]] = {} reasoning_content: str | None = None content_block: Final = ContentBlockDelta(**chunk) thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = [] @@ -809,8 +813,8 @@ class ModelResponseIterator: def _handle_redacted_thinking_content( self, content_block_start: ContentBlockStart, - provider_specific_fields: dict[str, Any], - ) -> tuple[list[ChatCompletionRedactedThinkingBlock], dict[str, Any]]: + provider_specific_fields: dict[str, object], + ) -> tuple[list[ChatCompletionRedactedThinkingBlock], dict[str, object]]: """ Handle the redacted thinking content """ @@ -878,7 +882,7 @@ class ModelResponseIterator: tool_use: ChatCompletionToolCallChunk | None = None finish_reason = "" usage: Usage | None = None - provider_specific_fields: dict[str, Any] = {} + provider_specific_fields: dict[str, object] = {} reasoning_content: str | None = None thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None @@ -1212,7 +1216,7 @@ class ModelResponseIterator: # Try to parse as valid JSON first try: - data_json: Final = json.loads(data_str) + data_json: Final = _loads_stream_chunk(data_str) return self.chunk_parser(chunk=data_json) except json.JSONDecodeError: # Switch to accumulation mode and start accumulating @@ -1330,7 +1334,7 @@ class ModelResponseIterator: str_line = str_line[index:] if str_line.startswith("data:"): - data_json: Final = json.loads(str_line[5:]) + data_json: Final = _loads_stream_chunk(str_line[5:]) return self.chunk_parser(chunk=data_json) else: return ModelResponseStream(id=self.response_id) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e1387a9068c..aa805ccea71 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, cast import httpx from pydantic import ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -125,7 +126,25 @@ else: _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 -_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[Any], bool]]] = MappingProxyType( + +class _AnthropicUsageIteration(TypedDict, total=False): + """One entry of the ``usage.iterations`` array on an Anthropic response.""" + + input_tokens: ReadOnly[int | None] + output_tokens: ReadOnly[int | None] + cache_creation_input_tokens: ReadOnly[int | None] + cache_read_input_tokens: ReadOnly[int | None] + + +class _AnthropicToolResultBlock(TypedDict, total=False): + """A ``*_tool_result`` content block on an Anthropic response.""" + + type: ReadOnly[str] + tool_use_id: ReadOnly[str] + content: ReadOnly[object] + + +_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[object], bool]]] = MappingProxyType( { "null": lambda v: v is None, "boolean": lambda v: isinstance(v, bool), @@ -440,7 +459,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("speed", None) @staticmethod - def _raise_invalid_reasoning_effort(model: str, value: Any, llm_provider: str) -> NoReturn: + def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn: """Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``. Args: @@ -1466,7 +1485,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) if _tool_choice is not None: - optional_params["tool_choice"] = _tool_choice + optional_params["tool_choice"] = AnthropicConfig._apply_forced_tool_choice( + model=model, tool_choice=_tool_choice, drop_params=drop_params + ) elif param == "stream" and value is True: optional_params["stream"] = value elif param == "stop" and (isinstance(value, str) or isinstance(value, list)): @@ -1495,7 +1516,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled) if _tool is None: continue - if not is_thinking_enabled: + if not is_thinking_enabled and not AnthropicModelInfo.forced_tool_use_unsupported(model): _tool_choice = { "name": RESPONSE_FORMAT_TOOL_NAME, "type": "tool", @@ -1992,19 +2013,35 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return data def _apply_output_config(self, data: dict, model: str, optional_params: dict) -> None: - """Validate and apply output_config to the request data.""" + """Validate and apply output_config to the request data. + + The ``drop_params`` gate here is an effort gate: ``format`` is a + structured-output field, not an effort field, so it survives the drop + and is vetted where it is consumed (the map's + ``supports_native_structured_output`` flag on emission paths). + """ if "output_config" not in optional_params: return output_config: Final = optional_params.get("output_config") if not output_config or not isinstance(output_config, dict): return - if litellm.drop_params is True and not self._model_supports_effort_param(model, self._resolved_provider): + if ( + litellm.drop_params is True + and any(key != "format" for key in output_config) + and not self._model_supports_effort_param(model, self._resolved_provider) + ): litellm.verbose_logger.warning( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, model, ) - optional_params.pop("output_config", None) - data.pop("output_config", None) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + optional_params.pop("output_config", None) + data.pop("output_config", None) + return + format_only: Final = {"format": preserved_format} # mutable-ok: json body + optional_params["output_config"] = format_only # rebind-ok: out-param store + data["output_config"] = format_only # rebind-ok: out-param store return effort: Final = output_config.get("effort") valid_efforts: Final = ["high", "medium", "low", "xhigh", "max"] @@ -2059,22 +2096,22 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self, completion_response: dict ) -> tuple[ str, - list[Any] | None, + list[object] | None, list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, str | None, list[ChatCompletionToolCallChunk], - list[Any] | None, - list[Any] | None, - list[Any] | None, + list[object] | None, + list[_AnthropicToolResultBlock] | None, + list[object] | None, ]: text_content = "" - citations: list[Any] | None = None + citations: list[object] | None = None thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None = None reasoning_content: str | None = None tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] - web_search_results: list[Any] | None = None - tool_results: list[Any] | None = None - compaction_blocks: list[Any] | None = None + web_search_results: list[object] | None = None + tool_results: list[_AnthropicToolResultBlock] | None = None + compaction_blocks: list[object] | None = None for idx, content in enumerate(completion_response["content"]): if content["type"] == "text": text_content += content["text"] @@ -2284,7 +2321,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): raw_speed: Final = _usage.get("speed") resolved_speed: Final = raw_speed if isinstance(raw_speed, str) else speed - iterations: Final[list[Any] | None] = _usage.get("iterations") + iterations: Final[Sequence[_AnthropicUsageIteration] | None] = _usage.get("iterations") if iterations: prompt_tokens = sum(it.get("input_tokens", 0) or 0 for it in iterations) completion_tokens = sum(it.get("output_tokens", 0) or 0 for it in iterations) @@ -2377,7 +2414,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _build_code_interpreter_results( self, - tool_results: list[Any], + tool_results: Sequence[_AnthropicToolResultBlock], code_by_id: dict[str, str], container_id: str | None, ) -> list[OutputCodeInterpreterCall]: @@ -2403,11 +2440,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _build_provider_specific_fields( self, completion_response: dict, - citations: list[Any] | None, + citations: Sequence[object] | None, thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] | None, - web_search_results: list[Any] | None, - tool_results: list[Any] | None, - compaction_blocks: list[Any] | None, + web_search_results: Sequence[object] | None, + tool_results: Sequence[_AnthropicToolResultBlock] | None, + compaction_blocks: Sequence[object] | None, tool_calls: list[ChatCompletionToolCallChunk], ) -> dict[str, Any]: provider_specific_fields: Final[dict[str, Any]] = { diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d8d6a7fc9f8..d23690976ad 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -28,10 +28,15 @@ from litellm.types.llms.anthropic import ( ANTHROPIC_OAUTH_TOKEN_PREFIX, AllAnthropicToolsValues, AnthropicMcpServerTool, + AnthropicMessagesToolChoice, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.model_listing import ModelInfoResponse +DROP_FORCED_TOOL_CHOICE_WARNING: Final = ( + "Downgrading forced tool_choice to 'auto' for model=%s (drop_params=True): this model rejects tool_choice type " + "'any'/'tool' with a 400 because thinking is always on and a forced call would skip it." +) DROP_DISABLED_THINKING_WARNING: Final = ( "Dropping `thinking={'type': 'disabled'}` for model=%s: thinking is always on for this model and cannot be " "disabled (the alternative is a provider 400). The model will still think adaptively, its response can contain " @@ -320,6 +325,45 @@ class AnthropicModelInfo(BaseLLMModelInfo): status_code=400, ) + @staticmethod + def forced_tool_use_unsupported(model: str) -> bool: + return AnthropicModelInfo._get_model_capability(model, "supports_forced_tool_use") is False + + @staticmethod + def forced_tool_use_downgraded(model: str, drop_params: bool) -> bool: + """True when the model map flags the model with + ``supports_forced_tool_use: false`` (Fable 5.1 / Mythos 5.1 400 on + ``any``/``tool``) and ``drop_params`` asks for the ``auto`` downgrade; + raises a clean client-side 400 for such models without ``drop_params``.""" + if not AnthropicModelInfo.forced_tool_use_unsupported(model): + return False + if not (litellm.drop_params or drop_params): + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} does not support forced tool use (tool_choice='required' or a named tool). " + "Use tool_choice='auto' and tell the model in the prompt when to call the tool, or set " + "`litellm.drop_params = True` to downgrade to 'auto' automatically." + ), + status_code=400, + ) + litellm.verbose_logger.warning(DROP_FORCED_TOOL_CHOICE_WARNING, model) + return True + + @staticmethod + def _apply_forced_tool_choice( + model: str, + tool_choice: AnthropicMessagesToolChoice, + drop_params: bool, + ) -> AnthropicMessagesToolChoice: + if tool_choice["type"] not in ("any", "tool"): + return tool_choice + if not AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params): + return tool_choice + disable_parallel: Final = tool_choice.get("disable_parallel_tool_use") + if disable_parallel is None: + return AnthropicMessagesToolChoice(type="auto") + return AnthropicMessagesToolChoice(type="auto", disable_parallel_tool_use=disable_parallel) + @staticmethod def _strip_version_suffix(model: str) -> str: at: Final = model.rfind("@") @@ -865,13 +909,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Anthropic. Status code: {response.status_code}, Response: {response.text}" ) - models: Final = response.json()["data"] + models: Final[Sequence[Mapping[str, str]]] = response.json()["data"] - litellm_model_names: Final = [] - for model in models: - stripped_model_name = model["id"] - litellm_model_name = "anthropic/" + stripped_model_name - litellm_model_names.append(litellm_model_name) + litellm_model_names: Final = ["anthropic/" + model["id"] for model in models] return litellm_model_names def get_token_counter(self) -> BaseTokenCounter | None: @@ -1077,7 +1117,7 @@ def strip_empty_content_blocks_from_anthropic_messages( return out -def _is_empty_text_block(block: Any) -> bool: +def _is_empty_text_block(block: object) -> bool: if not isinstance(block, dict) or block.get("type") != "text": return False text: Final = block.get("text") @@ -1131,7 +1171,7 @@ def normalize_anthropic_tool_use_id(raw_id: str) -> str: return sanitized or "tool_use_id" -def _sanitize_tool_use_id_content_block(block: Any) -> Any: +def _sanitize_tool_use_id_content_block(block: object) -> object: if not isinstance(block, dict): return block block_type: Final = block.get("type") @@ -1338,31 +1378,38 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: return additional_headers -def _anthropic_model_entry(model: ModelInfoResponse, created_at: str) -> Mapping[str, object]: +def _anthropic_model_entry( + model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str] +) -> Mapping[str, object]: return { # mutable-ok: JSON response body, serialized by the route and never mutated "type": "model", "id": model["id"], - "display_name": model["id"], + "display_name": display_names.get(model["id"], model["id"]), "created_at": created_at, "max_input_tokens": model.get("max_input_tokens"), "max_tokens": model.get("max_output_tokens"), } -def create_anthropic_model_list_response(models: Sequence[ModelInfoResponse]) -> Mapping[str, object]: +def create_anthropic_model_list_response( + models: Sequence[ModelInfoResponse], + display_names: Mapping[str, str] = MappingProxyType({}), +) -> Mapping[str, object]: """Build the Anthropic-native /v1/models envelope. Clients that send an anthropic-version header parse the Anthropic Models API shape (type/display_name/created_at plus has_more/first_id/last_id) and filter the list themselves, so every model is returned here. The token limits carry over from the OpenAI-shaped listing, named as the Messages API names them, and - are always present because the vendor shape declares them nullable, not optional + are always present because the vendor shape declares them nullable, not optional. + display_names maps a listed model id to a configured human-readable name; ids + without an entry fall back to the id itself, matching the vendor behavior """ created_at: Final = ( datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") ) data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated - _anthropic_model_entry(model, created_at) for model in models + _anthropic_model_entry(model, created_at, display_names) for model in models ] return { # mutable-ok: JSON response body, serialized by the route and never mutated "data": data, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 411df267442..199a8ab77e7 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -18,6 +18,24 @@ TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LE PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"}) +def _optional_attr(source: object, name: str) -> object: + return getattr(source, name, None) + + +def _as_string_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +def _thought_signature(provider_specific_fields: object) -> str | None: + fields: Final = _as_string_mapping(provider_specific_fields) + if fields is None: + return None + signature: Final = fields.get("thought_signature") + return signature if isinstance(signature, str) else None + + _ANTHROPIC_TOOL_SCHEMA_KEYS: Final = frozenset( {"name", "type", "input_schema", "description", "cache_control", "strict"} ) @@ -56,7 +74,7 @@ def truncate_tool_name(name: str) -> str: def create_tool_name_mapping( - tools: list[dict[str, Any]], + tools: Sequence[Mapping[str, object]], ) -> dict[str, str]: """ Create a mapping of truncated tool names to original names. @@ -70,6 +88,8 @@ def create_tool_name_mapping( mapping: Final[dict[str, str]] = {} for tool in tools: original_name = tool.get("name", "") + if not isinstance(original_name, str): + continue truncated_name = truncate_tool_name(original_name) if truncated_name != original_name: mapping[truncated_name] = original_name @@ -286,44 +306,44 @@ class LiteLLMAnthropicMessagesAdapter: ### FOR [BETA] `/v1/messages` endpoint support - def _extract_signature_from_tool_call(self, tool_call: Any) -> str | None: + def _extract_signature_from_tool_call(self, tool_call: object) -> str | None: """ Extract signature from a tool call's provider_specific_fields. Only checks provider_specific_fields, not thinking blocks. """ - signature = None + fields: Final = _optional_attr(tool_call, "provider_specific_fields") + if fields: + return _thought_signature(fields) - if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: - if "thought_signature" in tool_call.provider_specific_fields: - signature = tool_call.provider_specific_fields["thought_signature"] - elif hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields: - if "thought_signature" in tool_call.function.provider_specific_fields: - signature = tool_call.function.provider_specific_fields["thought_signature"] + function_fields: Final = _optional_attr(_optional_attr(tool_call, "function"), "provider_specific_fields") + if function_fields: + return _thought_signature(function_fields) - return signature + return None - def _extract_signature_from_tool_use_content(self, content: dict[str, Any]) -> str | None: + def _extract_signature_from_tool_use_content(self, content: Mapping[str, object]) -> str | None: """ Extract signature from a tool_use content block's provider_specific_fields. """ - provider_specific_fields: Final = content.get("provider_specific_fields", {}) + provider_specific_fields: Final = _as_string_mapping(content.get("provider_specific_fields", {})) if provider_specific_fields: - return provider_specific_fields.get("signature") + signature: Final = provider_specific_fields.get("signature") + return signature if isinstance(signature, str) else None return None def _add_cache_control_if_applicable( self, - source: Any, - target: Any, + source: object, + target: object, model: str | None, ) -> None: """ Extract cache_control from source and add to target if it should be preserved. - This method accepts Any type to support both regular dicts and TypedDict objects. - TypedDict objects (like ChatCompletionTextObject, ChatCompletionImageObject, etc.) - are dicts at runtime but have specific types at type-check time. Using Any allows - this method to work with both while maintaining runtime correctness. + This method accepts an unconstrained type to support both regular dicts and + TypedDict objects. TypedDict objects (like ChatCompletionTextObject, + ChatCompletionImageObject, etc.) are dicts at runtime but have specific types at + type-check time, so the widest parameter type works with both. Args: source: Dict or TypedDict containing potential cache_control field @@ -751,7 +771,7 @@ class LiteLLMAnthropicMessagesAdapter: return new_tools, tool_name_mapping - def translate_anthropic_output_format_to_openai(self, output_format: Any) -> dict[str, object] | None: + def translate_anthropic_output_format_to_openai(self, output_format: object) -> dict[str, object] | None: """ Translate Anthropic's output_format to OpenAI's response_format. @@ -1366,7 +1386,7 @@ class LiteLLMAnthropicMessagesAdapter: @classmethod def _first_positive_prompt_tokens_detail_value(cls, usage: Usage, field_names: tuple[str, ...]) -> int: - prompt_tokens_details: Final = getattr(usage, "prompt_tokens_details", None) + prompt_tokens_details: Final = _optional_attr(usage, "prompt_tokens_details") if prompt_tokens_details is None: return 0 @@ -1374,7 +1394,7 @@ class LiteLLMAnthropicMessagesAdapter: if isinstance(prompt_tokens_details, dict): value = cls._positive_int(prompt_tokens_details.get(field_name)) else: - value = cls._positive_int(getattr(prompt_tokens_details, field_name, None)) + value = cls._positive_int(_optional_attr(prompt_tokens_details, field_name)) if value > 0: return value return 0 diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 1cf52045a48..050ab67c86c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -14,7 +14,7 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: import re from collections.abc import Awaitable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, Union, cast from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack @@ -232,7 +232,7 @@ async def _check_summary_model_access( key_models: Final = list(getattr(user_api_key_auth, "models", None) or []) team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) - team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None) + team_model_aliases: Final[dict[str, str] | None] = getattr(user_api_key_auth, "team_model_aliases", None) team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or []) user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None) @@ -443,7 +443,9 @@ async def _check_summary_model_budget( ) return False - end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) + end_user_model_max_budget: Final[dict[str, object] | None] = getattr( + user_api_key_auth, "end_user_model_max_budget", None + ) end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: @@ -854,8 +856,8 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( - system: str | list[dict[str, Any]] | None, -) -> Mapping[str, object] | None: + system: str | list[dict[str, object]] | None, +) -> dict[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. Accepts a bare string or a list of Anthropic content blocks; returns @@ -866,10 +868,10 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts: Final[tuple[str, ...]] = tuple( + parts: Final[list[object]] = [ block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text" - ) - joined: Final = "\n\n".join(part for part in parts if part) + ] + joined: Final = "\n\n".join(part for part in parts if isinstance(part, str) and part) return {"role": "system", "content": joined} if joined else None return None @@ -951,7 +953,7 @@ async def _call_summary_model( summary_model: str, summary_messages: Sequence[Mapping[str, object]], metadata: Mapping[str, object], - llm_router: object, + llm_router: Optional["Router"], allowed_model_region: str | None = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, ) -> Union["ModelResponse", "CustomStreamWrapper"]: @@ -1036,10 +1038,9 @@ def _extract_usage(response: object) -> tuple[int, int]: usage: Final[object] = getattr(response, "usage", None) if usage is None: return 0, 0 - return ( - int(getattr(usage, "prompt_tokens", 0) or 0), - int(getattr(usage, "completion_tokens", 0) or 0), - ) + prompt_tokens: Final[int | None] = getattr(usage, "prompt_tokens", 0) + completion_tokens: Final[int | None] = getattr(usage, "completion_tokens", 0) + return int(prompt_tokens or 0), int(completion_tokens or 0) def apply_client_compaction_block_history( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 45c7825344b..66e36dab2ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None: return None -def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: +def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None: if isinstance(chunk, dict): - return chunk if chunk.get("type") == "error" else None + return chunk if chunk.get("type") == event_type else None if isinstance(chunk, (bytes, bytearray)): decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines()) return next( ( candidate for candidate in decoded_lines - if isinstance(candidate, dict) and candidate.get("type") == "error" + if isinstance(candidate, dict) and candidate.get("type") == event_type ), None, ) return None +def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: + return _anthropic_event_payload(chunk, "error") + + +def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None: + """ + Return the ``stop_details`` object of an Anthropic SSE ``message_delta`` + chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other chunk, a plain refusal without ``stop_details`` included. + """ + payload: Final = _anthropic_event_payload(chunk, "message_delta") + delta: Final = payload.get("delta") if payload is not None else None + if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal": + return None + stop_details: Final = delta.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None: """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None.""" payload: Final = _anthropic_error_event_payload(chunk) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 02d82887dde..9deff950724 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -1,11 +1,40 @@ +from collections.abc import Mapping from functools import lru_cache -from typing import Any, Final, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +if TYPE_CHECKING: + from litellm.exceptions import ContentPolicyViolationError + + +def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None: + """ + Return the ``stop_details`` of an Anthropic Messages response refused by a + safeguard (``stop_reason: "refusal"`` carrying ``stop_details``: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other response, a plain refusal without ``stop_details`` included. + """ + if not isinstance(response, dict) or response.get("stop_reason") != "refusal": + return None + stop_details: Final = response.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + +def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError": + """The exception a safeguard-refused Anthropic response converts into so the + content-policy fallback chain can re-dispatch it.""" + from litellm.exceptions import ContentPolicyViolationError + + return ContentPolicyViolationError( + message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).", + model=model, + llm_provider="anthropic", + ) + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: @@ -100,14 +129,12 @@ def mock_response( model=model, ) return AnthropicMessagesResponse( - **{ - "content": [{"text": mock_response, "type": "text"}], - "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", - "model": "claude-sonnet-4-20250514", - "role": "assistant", - "stop_reason": "end_turn", - "stop_sequence": None, - "type": "message", - "usage": {"input_tokens": 2095, "output_tokens": 503}, - } + content=[{"text": mock_response, "type": "text"}], + id="msg_013Zva2CMHLNnXjNJJKqJ2EF", + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage={"input_tokens": 2095, "output_tokens": 503}, ) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index ace7fc25dc9..0eb0e38a46e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -179,14 +179,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ) @staticmethod - def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, Any]]) -> str: + def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str: """Group a run of consecutive thinking blocks together; keep every other block alone.""" index, block = indexed_block return "thinking" if block.get("type") == "thinking" else f"block:{index}" @classmethod def _assistant_group_to_input_item( - cls, group: tuple[Mapping[str, Any], ...] + cls, group: tuple[Mapping[str, object], ...] ) -> dict[str, Any] | None: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") @@ -206,7 +206,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_messages_to_responses_input( self, messages: list[AllAnthropicPassThroughMessageValues], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Convert Anthropic messages list to Responses API `input` items. @@ -220,7 +220,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: assistant thinking -> reasoning assistant tool_use -> function_call """ - input_items: Final[list[dict[str, Any]]] = [] + input_items: Final[list[dict[str, object]]] = [] for m in messages: if m["role"] == "system": @@ -248,7 +248,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: } ) elif isinstance(content, list): - user_parts: list[dict[str, Any]] = [] + user_parts: list[Mapping[str, object]] = [] tool_image_parts: list[dict[str, Any]] = [] # mutable-ok: json content parts for block in content: if not isinstance(block, dict): @@ -379,9 +379,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_tools_to_responses_api( self, tools: list[AllAnthropicToolsValues], - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """Convert Anthropic tool definitions to Responses API function tools.""" - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in tools: tool_dict = cast(dict[str, Any], tool) tool_type = tool_dict.get("type", "") @@ -392,7 +392,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue # Responses turns strict mode on when `strict` is omitted, silently rewriting # `required` to every property. Anthropic tools are non-strict unless asked. - func_tool: dict[str, Any] = { + func_tool: dict[str, object] = { "type": "function", "name": tool_name, "strict": bool(tool_dict.get("strict")), @@ -407,7 +407,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_tool_choice_to_responses_api( tool_choice: AnthropicMessagesToolChoice, - ) -> str | dict[str, Any]: + ) -> str | dict[str, object]: """Convert Anthropic tool_choice to Responses API tool_choice.""" tc_type: Final = tool_choice.get("type") if tc_type == "any": @@ -420,8 +420,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_context_management_to_responses_api( - context_management: dict[str, Any], - ) -> list[dict[str, Any]] | None: + context_management: dict[str, object], + ) -> list[dict[str, object]] | None: """ Convert Anthropic context_management dict to OpenAI Responses API array format. @@ -435,13 +435,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if not isinstance(edits, list): return None - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for edit in edits: if not isinstance(edit, dict): continue edit_type = edit.get("type", "") if edit_type == "compact_20260112": - entry: dict[str, Any] = {"type": "compaction"} + entry: dict[str, object] = {"type": "compaction"} trigger = edit.get("trigger") if isinstance(trigger, dict) and trigger.get("value") is not None: entry["compact_threshold"] = int(trigger["value"]) @@ -451,9 +451,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @staticmethod def translate_thinking_to_reasoning( - thinking: dict[str, Any], - output_config: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + thinking: dict[str, object], + output_config: dict[str, object] | None = None, + ) -> dict[str, object] | None: """ Convert Anthropic thinking param to Responses API reasoning param. @@ -473,12 +473,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(output_config, dict) and output_config.get("effort"): effort = output_config["effort"] elif thinking_type == "enabled": - effort = reasoning_effort_from_thinking_budget(thinking.get("budget_tokens", 0)) + raw_budget: Final = thinking.get("budget_tokens", 0) + budget_tokens: Final = int(raw_budget) if isinstance(raw_budget, (int, float)) else 0 + effort = reasoning_effort_from_thinking_budget(budget_tokens) else: return None auto_summary: Final = is_reasoning_auto_summary_enabled() - result: Final[dict[str, Any]] = {"effort": effort} + result: Final[dict[str, object]] = {"effort": effort} summary: Final = thinking.get("summary") if summary: result["summary"] = summary @@ -570,7 +572,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # output_format / output_config.format -> text format # output_format: {"type": "json_schema", "schema": {...}} # output_config: {"format": {"type": "json_schema", "schema": {...}}} - output_format: Any = anthropic_request.get("output_format") + output_format: object = anthropic_request.get("output_format") output_config = anthropic_request.get("output_config") if not isinstance(output_format, dict) and isinstance(output_config, dict): output_format = output_config.get("format") @@ -620,7 +622,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ResponseReasoningItem, ) - content: Final[list[dict[str, Any]]] = [] + content: Final[list[dict[str, object]]] = [] stop_reason: AnthropicFinishReason = "end_turn" for item in response.output: diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 5fdf2ceff7f..dfd62ca575b 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Coroutine -from typing import Any, Final +from typing import Final import httpx @@ -116,7 +116,7 @@ class AnthropicFilesHandler: api_key: str | None = None, timeout: float | httpx.Timeout = 600.0, max_retries: int | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: """ Retrieve file content from Anthropic. diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 2bcc830851a..46a9dd1a531 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Any, Final +from typing import Final import httpx from openai import ( @@ -374,7 +374,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) error_body: Final = getattr(e, "body", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) @@ -392,7 +392,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): model: str, api_base: str, data: dict, - timeout: Any, + timeout: float | httpx.Timeout, dynamic_params: bool, model_response: ModelResponse, logging_obj: LiteLLMLoggingObj, @@ -502,7 +502,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): dynamic_params: bool, data: dict[str, object], model: str, - timeout: Any, + timeout: float | httpx.Timeout, max_retries: int, azure_ad_token: str | None = None, azure_ad_token_provider: Callable | None = None, @@ -578,7 +578,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): dynamic_params: bool, data: dict, model: str, - timeout: Any, + timeout: float | httpx.Timeout, max_retries: int, azure_ad_token: str | None = None, azure_ad_token_provider: Callable | None = None, @@ -634,7 +634,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except Exception as e: status_code: Final = getattr(e, "status_code", 500) error_headers = getattr(e, "headers", None) - error_response: Final = getattr(e, "response", None) + error_response: Final[object] = getattr(e, "response", None) message: Final = getattr(e, "message", str(e)) error_body: Final = getattr(e, "body", None) if error_headers is None and error_response: @@ -754,7 +754,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): aembedding=None, headers: dict | None = None, litellm_params: dict | None = None, - ) -> EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse]: + ) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]: if headers: optional_params["extra_headers"] = headers if self._client_session is None: @@ -1268,7 +1268,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers["Authorization"] = f"Bearer {azure_ad_token}" # init AzureOpenAI Client - azure_client_params: Final[dict[str, Any]] = self.initialize_azure_sdk_client( + azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client( litellm_params=litellm_params or {}, api_key=api_key, model_name=model or "", diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 2df4ab731ab..0ac0662205a 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from httpx._models import Headers, Response @@ -6,6 +8,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, hoist_images_from_tool_messages, + tool_with_flattened_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -32,6 +35,19 @@ else: LoggingClass = Any +_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: + tools: Final = optional_params.get("tools") + if not isinstance(tools, list): + return _NO_TOOLS_UPDATE + flattened: Final = [ # mutable-ok: request tools are a JSON list + tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + ] + return MappingProxyType({"tools": flattened}) + + class AzureOpenAIConfig(BaseConfig): """ Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions @@ -261,6 +277,7 @@ class AzureOpenAIConfig(BaseConfig): "model": model, "messages": azure_messages, **optional_params, + **flattened_tools_update(optional_params), } def transform_response( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 6cbd91bab5d..246bf69cb5f 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -20,6 +20,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_model_info, supports_reasoning from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig +from .gpt_transformation import flattened_tools_update class AzureOpenAIO1Config(OpenAIOSeriesConfig): @@ -108,4 +109,8 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): headers: dict, ) -> dict: model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name - return super().transform_request(model, messages, optional_params, litellm_params, headers) + flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict + **optional_params, + **flattened_tools_update(optional_params), + } + return super().transform_request(model, messages, flattened_params, litellm_params, headers) diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index a13b1300e55..f7382190fca 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -51,15 +51,13 @@ else: AsyncHTTPHandler = Any -class _AzureRawAnnotation(TypedDict, total=False): - type: ReadOnly[str] +class _AzureRawAnnotation(ChatCompletionAnnotation, total=False): text: ReadOnly[str] start_index: ReadOnly[int] end_index: ReadOnly[int] - url_citation: ReadOnly[ChatCompletionAnnotationURLCitation] -_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation | _AzureRawAnnotation +_TransformedAnnotation: TypeAlias = ChatCompletionAnnotation class _AzureText(TypedDict, total=False): @@ -223,18 +221,11 @@ class AzureAIAgentsHandler: """Build the ModelResponse from agent output.""" from litellm.types.utils import Choices, Message, Usage - message_kwargs: Final[dict[str, Any]] = { - "content": content, - "role": "assistant", - } - if annotations: - message_kwargs["annotations"] = annotations - model_response.choices = [ Choices( finish_reason="stop", index=0, - message=Message(**message_kwargs), + message=Message(content=content, role="assistant", annotations=annotations or None), ) ] model_response.model = model @@ -655,9 +646,6 @@ class AzureAIAgentsHandler: if data_str == "[DONE]": # Send final chunk with finish_reason - final_delta_kwargs: dict[str, Any] = {"content": None} - if collected_annotations: - final_delta_kwargs["annotations"] = collected_annotations final_chunk = ModelResponseStream( id=response_id, created=created, @@ -667,7 +655,7 @@ class AzureAIAgentsHandler: StreamingChoices( finish_reason="stop", index=0, - delta=Delta(**final_delta_kwargs), + delta=Delta(content=None, annotations=collected_annotations or None), ) ], ) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index b07a6b986d7..220fcedb0f8 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -155,8 +155,8 @@ class BaseTranslation(ABC): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[Any] | None = None, - ) -> list[bytes] | None: + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[bytes] | None: """ Build the streaming chunks that deliver a guardrail block message and cleanly terminate the stream in this provider's wire format. diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index f09ee210e6c..9b6f9c47105 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -124,6 +124,61 @@ def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage: ) +def stream_item_field(item: object, field: str) -> object | None: + if isinstance(item, dict): + return item.get(field) + return getattr(item, field, None) + + +def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: + """ + ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked + chat completions stream. + + A mid-stream block carries the chunks received so far as a list; real usage + rides on the final chunk when the upstream sent one + (``stream_options.include_usage``). Non-list originals defer to + ``blocked_response_usage``. + """ + if not isinstance(original_response, list): + usage: Final = blocked_response_usage(original_response) + return usage.get("input_tokens", 0), usage.get("output_tokens", 0) + usage_obj: Final = next( + ( + chunk_usage + for item in reversed(original_response) + if (chunk_usage := stream_item_field(item, "usage")) is not None + ), + None, + ) + return ( + _usage_tokens(usage_obj, "prompt_tokens", "input_tokens"), + _usage_tokens(usage_obj, "completion_tokens", "output_tokens"), + ) + + +def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsage: + """ + ``ResponseAPIUsage`` for a synthetic guardrail-blocked /v1/responses stream. + + A mid-stream block carries the events received so far as a list; real usage + rides on the ``response.completed`` event's response when the upstream sent + one. Non-list originals defer to ``blocked_responses_api_usage``. + """ + if not isinstance(original_response, list): + return blocked_responses_api_usage(original_response) + completed: Final = next( + ( + response + for item in reversed(original_response) + if stream_item_field(item, "type") == "response.completed" + and (response := stream_item_field(item, "response")) is not None + ), + None, + ) + return blocked_responses_api_usage(completed) + + def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool: per: Final = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None) if per is not None: diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index cced330d873..4fbc0ce51b0 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -5,7 +5,8 @@ import base64 import json from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar, cast +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Generic, Protocol, TypeVar, cast, runtime_checkable from litellm import verbose_logger from litellm.llms.base_llm.managed_resources.isolation import ( @@ -38,6 +39,30 @@ else: ResourceObjectType = TypeVar("ResourceObjectType") +@runtime_checkable +class _HasIdentifier(Protocol): + id: str + + +class _ManagedResourceRecord(Protocol[ResourceObjectType]): + unified_resource_id: str + resource_object: ResourceObjectType + + def model_dump(self) -> dict[str, object]: ... + + +class _ManagedResourceTable(Protocol[ResourceObjectType]): + async def create(self, *, data: Mapping[str, object]) -> object: ... + + async def find_first(self, *, where: Mapping[str, object]) -> _ManagedResourceRecord[ResourceObjectType] | None: ... + + async def find_many( + self, *, where: Mapping[str, object], take: int, order: Mapping[str, str] + ) -> list[_ManagedResourceRecord[ResourceObjectType]]: ... + + async def delete(self, *, where: Mapping[str, object]) -> object: ... + + class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ Base class for managing resources with target_model_names support. @@ -64,6 +89,9 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client + def _resource_table(self) -> _ManagedResourceTable[ResourceObjectType]: + return getattr(self.prisma_client.db, self.table_name) + # ============================================================================ # ABSTRACT METHODS # ============================================================================ @@ -137,7 +165,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): litellm_parent_otel_span: Span | None, model_mappings: dict[str, str], user_api_key_dict: UserAPIKeyAuth, - additional_db_fields: dict[str, Any] | None = None, + additional_db_fields: Mapping[str, object] | None = None, ) -> None: """ Store unified resource ID with model mappings in cache and database. @@ -153,7 +181,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): verbose_logger.info("Storing LiteLLM Managed %s with id=%s in cache", self.resource_type, unified_resource_id) # Prepare cache data - cache_data: Final = { + cache_data: Final[dict[str, object]] = { "unified_resource_id": unified_resource_id, "resource_object": resource_object, "model_mappings": model_mappings, @@ -176,7 +204,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) # Prepare database data - db_data: Final = { + db_data: Final[dict[str, object]] = { "unified_resource_id": unified_resource_id, "model_mappings": json.dumps(model_mappings), "flat_model_resource_ids": list(model_mappings.values()), @@ -205,7 +233,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): db_data.update(additional_db_fields) # Store in database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() result: Final = await table.create(data=db_data) verbose_logger.debug( @@ -240,7 +268,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): return result # Check database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() db_object: Final = await table.find_first(where={"unified_resource_id": unified_resource_id}) if db_object: @@ -264,7 +292,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): The deleted resource object or None if not found """ # Get old value from database - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() initial_value: Final = await table.find_first(where={"unified_resource_id": unified_resource_id}) if initial_value is None: @@ -515,7 +543,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): user_api_key_dict: UserAPIKeyAuth, limit: int | None = None, after: str | None = None, - additional_filters: dict[str, Any] | None = None, + additional_filters: Mapping[str, object] | None = None, ) -> dict[str, Any]: """ List resources created by a user. @@ -533,7 +561,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): if owner_filter is None: return build_list_page([]) - where_clause: Final[dict[str, Any]] = {**owner_filter} + where_clause: Final[dict[str, object]] = {**owner_filter} if after: where_clause["id"] = {"gt": after} @@ -544,14 +572,14 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Fetch resources fetch_limit: Final = limit or 20 - table: Final = getattr(self.prisma_client.db, self.table_name) + table: Final = self._resource_table() resources: Final = await table.find_many( where=where_clause, take=fetch_limit, order={"created_at": "desc"}, ) - resource_objects: Final[list[Any]] = [] + resource_objects: Final[list[object]] = [] for resource in resources: try: # Stop once we have enough @@ -559,12 +587,13 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): break # Parse resource object - resource_data = resource.resource_object - if isinstance(resource_data, str): - resource_data = json.loads(resource_data) + stored_resource = resource.resource_object + resource_data: object = ( + json.loads(stored_resource) if isinstance(stored_resource, str) else stored_resource + ) # Set unified ID - if hasattr(resource_data, "id"): + if isinstance(resource_data, _HasIdentifier): resource_data.id = resource.unified_resource_id elif isinstance(resource_data, dict): resource_data["id"] = resource.unified_resource_id diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index d1c77186ea8..3b302837032 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -75,6 +75,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" pages_processed: int | None = None + pages_processed_annotation: int | None = None credits: float | None = None doc_size_bytes: int | None = None diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 852cfaa24f2..1e634ced29b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1442,7 +1442,7 @@ class BaseAWSLLM: @tracer.wrap() def get_request_headers( self, - credentials: Credentials, + credentials: Credentials | None, aws_region_name: str, extra_headers: dict | None, endpoint_url: str, @@ -1469,9 +1469,13 @@ class BaseAWSLLM: try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + if credentials is None: + raise NoCredentialsError() + # Filter headers for AWS signature calculation # AWS SigV4 only includes specific headers in signature calculation aws_signature_headers: Final = self._filter_headers_for_aws_signature(headers) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index ca5f1298360..7d5f99ca893 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,4 +1,6 @@ import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -24,6 +26,22 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call +def _sigv4_principal(credentials: Credentials | None) -> Mapping[str, str]: + if credentials is None: + return MappingProxyType({}) + return MappingProxyType( + { + key: value + for key, value in ( + ("aws_access_key_id", credentials.access_key), + ("aws_secret_access_key", credentials.secret_key), + ("aws_session_token", credentials.token), + ) + if value is not None + } + ) + + def make_sync_call( client: HTTPHandler | None, api_base: str, @@ -95,7 +113,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers={}, client: AsyncHTTPHandler | None = None, @@ -167,7 +185,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers: dict = {}, client: AsyncHTTPHandler | None = None, @@ -331,7 +349,7 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - credentials: Final[Credentials] = self.get_credentials( + credentials: Final[Credentials | None] = self.get_credentials( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -368,19 +386,13 @@ class BedrockConverseLLM(BaseAWSLLM): # The Rust core owns the whole call for the subset it accepts. Ask # before transforming so whichever path runs emits pre_call once, and # hand down the credentials, region and endpoint this handler already - # resolved so both paths sign as the same principal. + # resolved so both paths sign as the same principal. Bearer-token auth + # resolves no SigV4 principal at all, and each path reads that token + # itself. rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy **optional_params, - **{ # mutable-ok: merged into its mutable parent above - key: value - for key, value in ( - ("aws_access_key_id", credentials.access_key), - ("aws_secret_access_key", credentials.secret_key), - ("aws_session_token", credentials.token), - ("aws_region_name", aws_region_name), - ) - if value is not None - }, + **_sigv4_principal(credentials), + "aws_region_name": aws_region_name, } serves_via_rust: Final = rust_chat_completions_accepts( model=model, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 395d99a4caa..5363c3c0366 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -87,6 +87,7 @@ from ..common_utils import ( BedrockError, BedrockModelInfo, bedrock_converse_supports_parallel_tool_use_config, + bedrock_model_accepts_cache_points, get_anthropic_beta_from_headers, get_bedrock_tool_name, is_bedrock_application_inference_profile_arn, @@ -588,6 +589,10 @@ class AmazonConverseConfig(BaseConfig): supported_params.append("context_management") return supported_params + @staticmethod + def _auto_tool_choice() -> ToolChoiceValuesBlock: + return ToolChoiceValuesBlock(auto={}) + def map_tool_choice_values( self, model: str, tool_choice: str | dict, drop_params: bool ) -> ToolChoiceValuesBlock | None: @@ -600,10 +605,14 @@ class AmazonConverseConfig(BaseConfig): status_code=400, ) elif tool_choice == "required": + if AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params): + return self._auto_tool_choice() return ToolChoiceValuesBlock(any={}) elif tool_choice == "auto": - return ToolChoiceValuesBlock(auto={}) + return self._auto_tool_choice() elif isinstance(tool_choice, dict): + if AnthropicModelInfo.forced_tool_use_downgraded(model, drop_params): + return self._auto_tool_choice() # only supported for anthropic + mistral models - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html specific_tool: Final = SpecificToolChoiceBlock( name=make_valid_bedrock_tool_name(tool_choice.get("function", {}).get("name", "")) @@ -1065,6 +1074,7 @@ class AmazonConverseConfig(BaseConfig): if ( litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider) and not is_thinking_enabled + and not AnthropicModelInfo.forced_tool_use_unsupported(model) ): optional_params["tool_choice"] = ToolChoiceValuesBlock( tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME) @@ -1140,7 +1150,7 @@ class AmazonConverseConfig(BaseConfig): model: str | None = None, ) -> SystemContentBlock | ContentBlock | None: cache_control: Final = message_block.get("cache_control", None) - if cache_control is None: + if cache_control is None or not bedrock_model_accepts_cache_points(model): return None cache_point: Final = self._build_cache_point_block(cache_control, model) @@ -1604,7 +1614,7 @@ class AmazonConverseConfig(BaseConfig): # Append cachePoint to tools if cache_control_injection_points has tool_config cache_injection_points: Final = additional_request_params.pop("cache_control_injection_points", None) - if cache_injection_points and len(bedrock_tools) > 0: + if cache_injection_points and len(bedrock_tools) > 0 and bedrock_model_accepts_cache_points(model): for point in cache_injection_points: if point.get("location") == "tool_config": cache_point = self._build_cache_point_block(point.get("control"), model) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 40b90014f3b..2a4c38e71ea 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers -from litellm.litellm_core_utils.litellm_logging import verbose_logger +from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -12,21 +12,21 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse -from litellm.utils import _supports_factory if TYPE_CHECKING: import tiktoken @@ -76,10 +76,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): drop_params: bool, ) -> dict: # Force tool-based structured outputs for Bedrock Invoke - # (similar to VertexAI fix in #19201) - # Bedrock Invoke doesn't support output_format parameter + # (similar to VertexAI fix in #19201) unless the model map advertises + # native structured output + from litellm.utils import supports_native_structured_output + original_model: Final = model - if "response_format" in non_default_params: + if "response_format" in non_default_params and not supports_native_structured_output( + model=model, custom_llm_provider="bedrock" + ): # Use a model name that forces tool-based approach model = "claude-3-sonnet-20240229" @@ -103,6 +107,16 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Restore original model name model = original_model + # The stub model hides the original model from the parent's forced-tool-use backstop + response_format_tool_choice: Final = optional_params.get("tool_choice") + if ( + "response_format" in non_default_params + and isinstance(response_format_tool_choice, dict) + and response_format_tool_choice.get("name") == RESPONSE_FORMAT_TOOL_NAME + and AnthropicModelInfo.forced_tool_use_unsupported(original_model) + ): + optional_params.pop("tool_choice") + return optional_params @staticmethod @@ -212,36 +226,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("stream_chunk_size", None) - output_format: Final = anthropic_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_request, - ) - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_request, + ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_request, + ) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 72e3cc1b326..66ee5f10679 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -177,6 +177,95 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages +def _bedrock_model_supports(model: str, key: str) -> bool: + from litellm.utils import _supports_factory + + return _supports_factory(model=model, custom_llm_provider="bedrock", key=key) + + +def apply_bedrock_invoke_structured_output( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Route Anthropic structured-output params to what the Bedrock model supports. + + Consumes the legacy top-level ``output_format`` and the newer + ``output_config.format``, keeping the pre-existing precedence of the legacy + field when a request carries both. Models flagged + ``supports_native_structured_output`` in the model map get the schema + forwarded as ``output_config.format``, which Bedrock relays to the model for + enforced structured output. For every other model the schema is inlined into + the last user message as best-effort text, with a warning because nothing + enforces it. + """ + legacy_output_format: Final = request_body.pop("output_format", None) + output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body) + schema_format: Final = legacy_output_format if isinstance(legacy_output_format, dict) else output_config_format + if schema_format is None: + return + + if _bedrock_model_supports(model, "supports_native_structured_output"): + existing_output_config: Final = request_body.get("output_config") + if isinstance(existing_output_config, dict): + existing_output_config["format"] = schema_format + else: + request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param # mutable-ok: json + return + + verbose_logger.warning( + "Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` " + "in model_prices_and_context_window.json, so the JSON schema was inlined into " + "the last user message and is NOT enforced by the model.", + model, + ) + convert_bedrock_invoke_output_format_to_inline_schema( + output_format=schema_format, + request_body=request_body, + ) + + +def strip_unsupported_bedrock_invoke_output_config_keys( + model: str, + request_body: dict[str, object], # mutable-ok: edited in place like siblings +) -> None: + """ + Drop ``output_config`` keys the Bedrock model does not accept. + + ``format`` survives unconditionally: it is only attached for models whose map + entry advertises ``supports_native_structured_output``. Effort-bearing keys + survive only when the map flags ``supports_output_config`` or a + ``supports_*_reasoning_effort`` tier; otherwise they are dropped with a + warning so Bedrock does not reject the request. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + output_config: Final = request_body.get("output_config") + if not isinstance(output_config, dict): + return + if all(key == "format" for key in output_config): + return + if _bedrock_model_supports(model, "supports_output_config") or AnthropicConfig._model_supports_effort_param( + model, "bedrock" + ): + return + + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` keys for " + "model=%s: neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + preserved_format: Final = output_config.get("format") + if preserved_format is None: + request_body.pop("output_config", None) + else: + request_body["output_config"] = {"format": preserved_format} # rebind-ok: out-param # mutable-ok: json + + def normalize_custom_field_on_tools(request_body: dict) -> None: """ Drop the ``custom`` field from each tool, first hoisting a boolean @@ -727,6 +816,30 @@ def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: ) +def bedrock_model_accepts_cache_points(model: str | None) -> bool: + """ + Whether Converse ``cachePoint`` blocks may be sent to this model. + + Bedrock rejects requests carrying cachePoint blocks for models without prompt + caching support ("You invoked an unsupported model or your request did not allow + prompt caching"), so a model whose cost-map entry does not declare + ``supports_prompt_caching`` must not receive them. A model absent from the map + (an application inference profile ARN, a model newer than the map) keeps emitting + so existing caching setups never silently degrade. ``litellm.utils.supports_prompt_caching`` + is not reusable here: it returns False for unmapped models, the opposite polarity. + """ + if model is None: + return True + entries: Final = tuple( + entry + for candidate in (model, get_bedrock_base_model(model)) + if (entry := litellm.model_cost.get(candidate)) is not None + ) + if not entries: + return True + return any(entry.get("supports_prompt_caching") is True for entry in entries) + + def is_claude_4_5_on_bedrock(model: str) -> bool: """ Check if the model supports Bedrock prompt caching with an extended '1h' TTL @@ -1487,6 +1600,7 @@ class CommonBatchFilesUtils: aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Prepare the request data diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 13718d41cc1..e74c3802d20 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -113,6 +113,7 @@ class BedrockFilesHandler(BaseAWSLLM): aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Create S3 client diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index f442608a288..33b27943ad8 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -146,6 +146,7 @@ class _BedrockS3RequestParams(BaseModel): aws_role_name: str | None = None aws_web_identity_token: str | None = None aws_sts_endpoint: str | None = None + aws_external_id: str | None = None s3_region_name: str | None = None s3_endpoint_url: str | None = None @@ -1029,6 +1030,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_role_name=optional_params.get("aws_role_name"), aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + aws_external_id=optional_params.get("aws_external_id"), ) # Calculate SHA256 hash of the content (REQUIRED for S3) @@ -1290,6 +1292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_role_name=request_params.aws_role_name, aws_web_identity_token=request_params.aws_web_identity_token, aws_sts_endpoint=request_params.aws_sts_endpoint, + aws_external_id=request_params.aws_external_id, ) empty_body_hash: Final = hashlib.sha256(b"").hexdigest() diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index f74a290d773..6ff9f0155f9 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,14 +29,14 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( - convert_bedrock_invoke_output_format_to_inline_schema, + apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, - pop_bedrock_invoke_output_config_format, + strip_unsupported_bedrock_invoke_output_config_keys, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -51,7 +51,6 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk, ModelResponseStream from litellm.types.utils import GenericStreamingChunk as GChunk -from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -708,52 +707,25 @@ class AmazonAnthropicClaudeMessagesConfig( # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models) self._remove_ttl_from_cache_control(anthropic_messages_request=anthropic_messages_request, model=model) - # 5. Convert structured-output params to inline schema. - # Bedrock Invoke doesn't support top-level `output_format`; its - # accepted `output_config` subset is also narrower than Anthropic's, so - # consume the newer `output_config.format` shape here instead of - # forwarding it as an unknown nested key. + # 5. Route structured-output params (`output_format` / + # `output_config.format`) to native enforcement or the inline-schema + # fallback, then strip `output_config` keys the model does not accept. + # Ref: https://github.com/BerriAI/litellm/issues/22797 existing_output_config: Final = anthropic_messages_request.get("output_config") if isinstance(existing_output_config, dict): anthropic_messages_request["output_config"] = dict(existing_output_config) - output_format: Final = anthropic_messages_request.pop("output_format", None) - output_config_format: Final = pop_bedrock_invoke_output_config_format(anthropic_messages_request) - if output_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_format, - request_body=anthropic_messages_request, - ) - elif output_config_format: - convert_bedrock_invoke_output_format_to_inline_schema( - output_format=output_config_format, - request_body=anthropic_messages_request, - ) + apply_bedrock_invoke_structured_output( + model=model, + request_body=anthropic_messages_request, + ) normalize_bedrock_opus_output_config_effort( model=model, output_config=anthropic_messages_request.get("output_config"), ) - - # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, - # but older models do not — strip it to avoid request rejection. - # Ref: https://github.com/BerriAI/litellm/issues/22797 - if not ( - _supports_factory( - model=model, - custom_llm_provider="bedrock", - key="supports_output_config", - ) - or AnthropicConfig._model_supports_effort_param(model, "bedrock") - ): - if anthropic_messages_request.pop("output_config", None) is not None: - verbose_logger.warning( - "Bedrock Invoke: stripping unsupported `output_config` for " - "model=%s — neither `supports_output_config` nor any " - "`supports_*_reasoning_effort` flag is set in " - "model_prices_and_context_window.json. Add the capability " - "flag to the model JSON entry if this model accepts " - "`output_config`.", - model, - ) + strip_unsupported_bedrock_invoke_output_config_keys( + model=model, + request_body=anthropic_messages_request, + ) # 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) # Ref: https://github.com/BerriAI/litellm/issues/22847 @@ -774,9 +746,11 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True - and "output_config" in anthropic_messages_request + and isinstance(remaining_output_config, dict) + and any(key != "format" for key in remaining_output_config) and not AnthropicConfig._model_supports_effort_param(model, "bedrock") ): verbose_logger.warning( diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 28c2e446d10..1f4c81d6491 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -7,7 +7,7 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. import base64 import json import uuid as uuid_lib -from typing import Any, Final, cast +from typing import Final, cast from pydantic import BaseModel @@ -633,7 +633,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final = json.loads(message) + json_message: Final[dict[str, object]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -1182,7 +1182,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Create a function call arguments done event # This is a custom event format that matches what clients expect - function_call_event: Final[dict[str, Any]] = { + function_call_event: Final[dict[str, object]] = { "type": "response.function_call_arguments.done", "event_id": f"event_{uuid.uuid4()}", "response_id": current_response_id, diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index 1ff02a6f8d9..178acb0de0d 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -8,9 +8,11 @@ then we poll until the result is ready. import asyncio import time -from typing import Any, Final +from collections.abc import Coroutine, Mapping +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -33,6 +35,42 @@ from ..common_utils import ( from .transformation import BlackForestLabsImageEditConfig +class _BFLSubmitBody(TypedDict, total=False): + """Decoded body of the BFL submit response, which hands back a polling URL.""" + + errors: ReadOnly[object] + polling_url: ReadOnly[str] + + +class _BFLPollBody(TypedDict, total=False): + """Decoded body of a BFL polling response.""" + + status: ReadOnly[str] + + +class _BFLSubmitResponse(Protocol): + """The submit call's HTTP response, read for its status, body text and decoded body.""" + + @property + def status_code(self) -> int: ... + + @property + def text(self) -> str: ... + + def json(self) -> _BFLSubmitBody: ... + + +class _BFLPollResponse(Protocol): + """A polling call's HTTP response, read only for the task status it carries.""" + + def json(self) -> _BFLPollBody: ... + + +def _poll_status(response: _BFLPollResponse) -> str | None: + """Read the task status out of a BFL polling response body.""" + return response.json().get("status") + + class BlackForestLabsImageEdit: """ Black Forest Labs Image Edit handler. @@ -53,10 +91,10 @@ class BlackForestLabsImageEdit: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, aimage_edit: bool = False, - ) -> ImageResponse | Any: + ) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Main entry point for image edit requests. @@ -185,7 +223,7 @@ class BlackForestLabsImageEdit: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: """ @@ -281,7 +319,7 @@ class BlackForestLabsImageEdit: def _poll_for_result_sync( self, - initial_response: httpx.Response, + initial_response: _BFLSubmitResponse, headers: dict, sync_client: HTTPHandler, max_wait: float = DEFAULT_MAX_POLLING_TIME, @@ -356,8 +394,7 @@ class BlackForestLabsImageEdit: message=f"Polling failed: {response.text}", ) - data = response.json() - status = data.get("status") + status = _poll_status(response) verbose_logger.debug("BFL poll status: %s", status) @@ -383,7 +420,7 @@ class BlackForestLabsImageEdit: async def _poll_for_result_async( self, - initial_response: httpx.Response, + initial_response: _BFLSubmitResponse, headers: dict, async_client: AsyncHTTPHandler, max_wait: float = DEFAULT_MAX_POLLING_TIME, @@ -447,8 +484,7 @@ class BlackForestLabsImageEdit: message=f"Polling failed: {response.text}", ) - data = response.json() - status = data.get("status") + status = _poll_status(response) verbose_logger.debug("BFL poll status: %s", status) diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index 03e4999c5aa..879bef37b58 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -8,9 +8,11 @@ then we poll until the result is ready. import asyncio import time -from typing import Any, Final +from collections.abc import Coroutine, Mapping +from typing import Final, Protocol, TypedDict import httpx +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger @@ -33,6 +35,23 @@ from ..common_utils import ( from .transformation import BlackForestLabsImageGenerationConfig +class _BFLTaskPayload(TypedDict, total=False): + """The body BFL returns for a submitted or polled generation task.""" + + errors: ReadOnly[object] + polling_url: ReadOnly[str] + status: ReadOnly[str] + + +class _TaskJsonResponse(Protocol): + def json(self) -> _BFLTaskPayload: ... + + +def _task_payload(response: _TaskJsonResponse) -> _BFLTaskPayload: + """The JSON body of a BFL task submission or poll response.""" + return response.json() + + class BlackForestLabsImageGeneration: """ Black Forest Labs Image Generation handler. @@ -53,10 +72,10 @@ class BlackForestLabsImageGeneration: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, str] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, aimg_generation: bool = False, - ) -> ImageResponse | Any: + ) -> ImageResponse | Coroutine[object, object, ImageResponse]: """ Main entry point for image generation requests. @@ -187,7 +206,7 @@ class BlackForestLabsImageGeneration: litellm_params: GenericLiteLLMParams | dict, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout | None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, str] | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: """ @@ -305,7 +324,7 @@ class BlackForestLabsImageGeneration: # Parse initial response to get polling URL try: - response_data: Final = initial_response.json() + response_data: Final = _task_payload(initial_response) except Exception as e: raise BlackForestLabsError( status_code=initial_response.status_code, @@ -350,7 +369,7 @@ class BlackForestLabsImageGeneration: message=f"Polling failed: {response.text}", ) - data = response.json() + data = _task_payload(response) status = data.get("status") verbose_logger.debug("BFL poll status: %s", status) @@ -396,7 +415,7 @@ class BlackForestLabsImageGeneration: # Parse initial response to get polling URL try: - response_data: Final = initial_response.json() + response_data: Final = _task_payload(initial_response) except Exception as e: raise BlackForestLabsError( status_code=initial_response.status_code, @@ -441,7 +460,7 @@ class BlackForestLabsImageGeneration: message=f"Polling failed: {response.text}", ) - data = response.json() + data = _task_payload(response) status = data.get("status") verbose_logger.debug("BFL poll status: %s", status) diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 8c08b2bc33c..f8486d3b274 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -4,9 +4,10 @@ import json from collections.abc import Callable from functools import partial -from typing import Final +from typing import Final, Protocol import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -23,6 +24,53 @@ from litellm.types.utils import TextChoices from litellm.utils import CustomStreamWrapper, TextCompletionResponse +class _CodestralChoiceMessage(TypedDict): + """`choices[].message` of a Codestral FIM completion.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[str | None]] + + +class _CodestralChoice(TypedDict): + """One entry of `choices` in a Codestral FIM completion.""" + + index: ReadOnly[int] + message: ReadOnly[NotRequired[_CodestralChoiceMessage]] + finish_reason: ReadOnly[NotRequired[str | None]] + logprobs: ReadOnly[NotRequired[dict[str, object] | None]] + + +class _CodestralUsage(TypedDict): + """Token accounting returned alongside a Codestral FIM completion.""" + + prompt_tokens: ReadOnly[NotRequired[int]] + completion_tokens: ReadOnly[NotRequired[int]] + total_tokens: ReadOnly[NotRequired[int]] + + +class _CodestralCompletionResponse(TypedDict): + """Body returned by the Codestral `/v1/fim/completions` endpoint.""" + + id: ReadOnly[NotRequired[str]] + created: ReadOnly[NotRequired[int]] + model: ReadOnly[NotRequired[str]] + object: ReadOnly[NotRequired[str]] + usage: ReadOnly[NotRequired[_CodestralUsage]] + choices: ReadOnly[NotRequired[list[_CodestralChoice]]] + + +class _CodestralHTTPResponse(Protocol): + """The Codestral completion response as this handler reads it.""" + + @property + def status_code(self) -> int: ... + + @property + def text(self) -> str: ... + + def json(self) -> _CodestralCompletionResponse: ... + + class TextCompletionCodestralError(Exception): def __init__( self, @@ -115,7 +163,7 @@ class CodestralTextCompletion: def process_text_completion_response( self, model: str, - response: httpx.Response, + response: _CodestralHTTPResponse, model_response: TextCompletionResponse, stream: bool, logging_obj: LiteLLMLogging, diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b6586481fd3..73adf9c7455 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -3,6 +3,7 @@ import concurrent.futures import contextlib import os import ssl +import sys import typing import urllib.request from collections.abc import Callable, Generator @@ -75,10 +76,22 @@ except ImportError: pass +def _current_task_is_cancelling() -> bool: + task: Final = asyncio.current_task() + if task is None or sys.version_info < (3, 11): + return True + return task.cancelling() > 0 + + @contextlib.contextmanager def map_aiohttp_exceptions() -> Generator[None, None, None]: try: yield + except asyncio.CancelledError as exc: + # a closing connector cancels its shielded DNS task; that surfaces here without the request task being cancelled + if _current_task_is_cancelling(): + raise + raise httpx.ConnectError("aiohttp transport cancelled the request internally") from exc except Exception as exc: mapped_exc: type[Exception] | None = None diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 5ab7fbf3658..26e60fa959d 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -54,6 +54,9 @@ class DashScopeChatConfig(OpenAIGPTConfig): dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY") return api_base, dynamic_api_key + def _resolve_chat_api_base(self, api_base: str | None) -> str: + return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1" + def get_complete_url( self, api_base: str | None, @@ -66,10 +69,7 @@ class DashScopeChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default DashScope /chat/completions endpoint. """ - if not api_base: - api_base = "https://dashscope.aliyuncs.com/compatible-mode/v1" - - if not api_base.endswith("/chat/completions"): - api_base = f"{api_base}/chat/completions" - - return api_base + resolved_api_base: Final = self._resolve_chat_api_base(api_base) + if resolved_api_base.endswith("/chat/completions"): + return resolved_api_base + return f"{resolved_api_base}/chat/completions" diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index 9a7dd4da8d3..b7c97893a15 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -2,9 +2,89 @@ Common utilities for the DashScope LLM provider. """ +from typing import TYPE_CHECKING + import httpx from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, + ) + from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig + + +def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig": + if custom_llm_provider == "qwencloud": + from litellm.llms.dashscope.qwencloud import QwenCloudEmbeddingConfig + + return QwenCloudEmbeddingConfig() + if custom_llm_provider == "qwen_ai_platform": + from litellm.llms.dashscope.qwen_ai_platform import ( + QwenAIPlatformEmbeddingConfig, + ) + + return QwenAIPlatformEmbeddingConfig() + from litellm.llms.dashscope.embed.transformation import DashScopeEmbeddingConfig + + return DashScopeEmbeddingConfig() + + +def get_dashscope_family_rerank_config(custom_llm_provider: str) -> "BaseRerankConfig": + if custom_llm_provider == "qwencloud": + from litellm.llms.dashscope.qwencloud import QwenCloudRerankConfig + + return QwenCloudRerankConfig() + if custom_llm_provider == "qwen_ai_platform": + from litellm.llms.dashscope.qwen_ai_platform import QwenAIPlatformRerankConfig + + return QwenAIPlatformRerankConfig() + from litellm.llms.dashscope.rerank.transformation import DashScopeRerankConfig + + return DashScopeRerankConfig() + + +def get_dashscope_family_image_generation_config( + custom_llm_provider: str, +) -> "BaseImageGenerationConfig": + if custom_llm_provider == "qwencloud": + from litellm.llms.dashscope.qwencloud import QwenCloudImageGenerationConfig + + return QwenCloudImageGenerationConfig() + if custom_llm_provider == "qwen_ai_platform": + from litellm.llms.dashscope.qwen_ai_platform import ( + QwenAIPlatformImageGenerationConfig, + ) + + return QwenAIPlatformImageGenerationConfig() + from litellm.llms.dashscope.image_generation.transformation import ( + DashScopeImageGenerationConfig, + ) + + return DashScopeImageGenerationConfig() + + +def resolve_dashscope_family_api_key(custom_llm_provider: str, api_key: str | None) -> str | None: + if custom_llm_provider == "dashscope": + return api_key or get_secret_str("DASHSCOPE_API_KEY") + return api_key or get_secret_str(f"{custom_llm_provider.upper()}_API_KEY") or get_secret_str("DASHSCOPE_API_KEY") + + +def missing_dashscope_family_key_message(custom_llm_provider: str) -> str: + if custom_llm_provider == "qwencloud": + return ( + "Missing API key for QwenCloud. Set QWENCLOUD_API_KEY or " + "DASHSCOPE_API_KEY environment variable or pass api_key parameter." + ) + if custom_llm_provider == "qwen_ai_platform": + return ( + "Missing API key for Qwen AI Platform. Set QWEN_AI_PLATFORM_API_KEY or " + "DASHSCOPE_API_KEY environment variable or pass api_key parameter." + ) + return "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." class DashScopeError(BaseLLMException): diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 771ce140f66..dd5bee1fe8b 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -110,7 +110,7 @@ def _calculate_completion_cost( return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) -def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: +def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]: """ Calculate cost per token for Dashscope models. @@ -119,11 +119,12 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Args: model: Model name without provider prefix usage: LiteLLM Usage block + custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases Returns: Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd) """ - model_info: Final = get_model_info(model=model, custom_llm_provider="dashscope") + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) breakdown: Final = _extract_token_breakdown(usage) raw_tiers: Final = model_info.get("tiered_pricing") tiered_pricing: Final = raw_tiers if isinstance(raw_tiers, list) else None diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 6d13f1e53f7..63ee984a65c 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -62,6 +62,17 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): # for drop_params=False before this method is called. return optional_params + def _resolve_api_key(self, api_key: str | None) -> str: + resolved_api_key: Final = api_key if api_key is not None else get_secret_str("DASHSCOPE_API_KEY") + if resolved_api_key is None: + raise ValueError( + "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." + ) + return resolved_api_key + + def _resolve_embedding_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE + def validate_environment( self, headers: dict, @@ -72,17 +83,11 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): api_key: str | None = None, api_base: str | None = None, ) -> dict: - if api_key is None: - api_key = get_secret_str("DASHSCOPE_API_KEY") - if api_key is None: - raise ValueError( - "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." - ) - default_headers: Final = { + return { "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {self._resolve_api_key(api_key)}", + **headers, } - return {**default_headers, **headers} def get_complete_url( self, @@ -93,8 +98,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - base = api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE - base = base.rstrip("/") + base: Final = self._resolve_embedding_api_base(api_base).rstrip("/") if base.endswith("/embeddings"): return base return f"{base}/embeddings" diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index a7f0e98865f..c0e278a96ef 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -91,6 +91,15 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): mapped[k] = v return mapped + def _resolve_api_key(self, api_key: str | None) -> str: + resolved_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY") + if not resolved_api_key: + raise ValueError("DASHSCOPE_API_KEY is not set") + return resolved_api_key + + def _resolve_image_api_base(self, image_api_base: str | None) -> str: + return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE + def get_complete_url( self, api_base: str | None, @@ -103,7 +112,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): image_api_base: Final = ( api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None ) - return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE + return self._resolve_image_api_base(image_api_base) def validate_environment( self, @@ -115,10 +124,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, api_base: str | None = None, ) -> dict: - final_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY") - if not final_api_key: - raise ValueError("DASHSCOPE_API_KEY is not set") - headers["Authorization"] = f"Bearer {final_api_key}" + headers["Authorization"] = f"Bearer {self._resolve_api_key(api_key)}" headers["Content-Type"] = "application/json" return headers diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py new file mode 100644 index 00000000000..9a44eaf574a --- /dev/null +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -0,0 +1,62 @@ +from typing import Final + +from litellm.secret_managers.main import get_secret_str + +from .chat.transformation import DashScopeChatConfig +from .embed.transformation import DashScopeEmbeddingConfig +from .image_generation.transformation import DashScopeImageGenerationConfig +from .rerank.transformation import DashScopeRerankConfig + +QWEN_AI_PLATFORM_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +QWEN_AI_PLATFORM_RERANK_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" +QWEN_AI_PLATFORM_IMAGE_API_BASE: Final = ( + "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" +) + + +def _resolve_qwen_ai_platform_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("QWEN_AI_PLATFORM_API_KEY") or get_secret_str("DASHSCOPE_API_KEY") + + +def _require_qwen_ai_platform_api_key(api_key: str | None) -> str: + resolved: Final = _resolve_qwen_ai_platform_api_key(api_key) + if resolved is None: + raise ValueError( + "Qwen AI Platform API key is required. Set 'QWEN_AI_PLATFORM_API_KEY' or 'DASHSCOPE_API_KEY' env var " + "or pass api_key explicitly." + ) + return resolved + + +class QwenAIPlatformChatConfig(DashScopeChatConfig): + def _get_openai_compatible_provider_info( + self, api_base: str | None, api_key: str | None + ) -> tuple[str | None, str | None]: + return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key) + + def _resolve_chat_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE + + +class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwen_ai_platform_api_key(api_key) + + def _resolve_embedding_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE + + +class QwenAIPlatformRerankConfig(DashScopeRerankConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwen_ai_platform_api_key(api_key) + + def _resolve_rerank_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_RERANK") or QWEN_AI_PLATFORM_RERANK_API_BASE + + +class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwen_ai_platform_api_key(api_key) + + def _resolve_image_api_base(self, image_api_base: str | None) -> str: + return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE diff --git a/litellm/llms/dashscope/qwencloud.py b/litellm/llms/dashscope/qwencloud.py new file mode 100644 index 00000000000..d8d53e340ef --- /dev/null +++ b/litellm/llms/dashscope/qwencloud.py @@ -0,0 +1,62 @@ +from typing import Final + +from litellm.secret_managers.main import get_secret_str + +from .chat.transformation import DashScopeChatConfig +from .embed.transformation import DashScopeEmbeddingConfig +from .image_generation.transformation import DashScopeImageGenerationConfig +from .rerank.transformation import DashScopeRerankConfig + +QWENCLOUD_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" +QWENCLOUD_RERANK_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks" +QWENCLOUD_IMAGE_API_BASE: Final = ( + "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" +) + + +def _resolve_qwencloud_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("QWENCLOUD_API_KEY") or get_secret_str("DASHSCOPE_API_KEY") + + +def _require_qwencloud_api_key(api_key: str | None) -> str: + resolved: Final = _resolve_qwencloud_api_key(api_key) + if resolved is None: + raise ValueError( + "QwenCloud API key is required. Set 'QWENCLOUD_API_KEY' or 'DASHSCOPE_API_KEY' env var " + "or pass api_key explicitly." + ) + return resolved + + +class QwenCloudChatConfig(DashScopeChatConfig): + def _get_openai_compatible_provider_info( + self, api_base: str | None, api_key: str | None + ) -> tuple[str | None, str | None]: + return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key) + + def _resolve_chat_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE + + +class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwencloud_api_key(api_key) + + def _resolve_embedding_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE + + +class QwenCloudRerankConfig(DashScopeRerankConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwencloud_api_key(api_key) + + def _resolve_rerank_api_base(self, api_base: str | None) -> str: + return api_base or get_secret_str("QWENCLOUD_API_BASE_RERANK") or QWENCLOUD_RERANK_API_BASE + + +class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig): + def _resolve_api_key(self, api_key: str | None) -> str: + return _require_qwencloud_api_key(api_key) + + def _resolve_image_api_base(self, image_api_base: str | None) -> str: + return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 98be4e4f2e7..3dd3996b2ee 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -58,19 +58,30 @@ class DashScopeRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass + def _resolve_api_key(self, api_key: str | None) -> str: + resolved_api_key: Final = api_key if api_key is not None else get_secret_str("DASHSCOPE_API_KEY") + if resolved_api_key is None: + raise ValueError( + "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." + ) + return resolved_api_key + + def _resolve_rerank_api_base(self, api_base: str | None) -> str: + if api_base is not None: + return api_base + return get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL + def get_complete_url( self, api_base: str | None, model: str, optional_params: dict | None = None, ) -> str: - if api_base is None: - api_base = get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL + resolved_api_base: Final = self._resolve_rerank_api_base(api_base) + if resolved_api_base == DEFAULT_RERANK_URL: + return resolved_api_base - if api_base == DEFAULT_RERANK_URL: - return DEFAULT_RERANK_URL - - cleaned: Final = api_base.rstrip("/") + cleaned: Final = resolved_api_base.rstrip("/") if cleaned.endswith("/reranks") or cleaned.endswith("/rerank"): return cleaned @@ -88,19 +99,12 @@ class DashScopeRerankConfig(BaseRerankConfig): optional_params: dict | None = None, litellm_params: Mapping[str, object] | None = None, ) -> dict: - if api_key is None: - api_key = get_secret_str("DASHSCOPE_API_KEY") - if api_key is None: - raise ValueError( - "DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly." - ) - - default_headers: Final = { - "Authorization": f"Bearer {api_key}", + return { + "Authorization": f"Bearer {self._resolve_api_key(api_key)}", "accept": "application/json", "content-type": "application/json", + **headers, } - return {**default_headers, **headers} def get_supported_cohere_rerank_params(self, model: str) -> list: return ["query", "documents", "top_n", "return_documents"] diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index e52c56af82b..a3d0482af0a 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -2,10 +2,11 @@ Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -24,6 +25,36 @@ from litellm.types.rerank import ( ) +class _DeepinfraInferenceStatus(TypedDict, total=False): + """The ``inference_status`` block of a DeepInfra rerank response.""" + + status: ReadOnly[str] + runtime_ms: ReadOnly[float] + cost: ReadOnly[float] + tokens_generated: ReadOnly[int] + tokens_input: ReadOnly[int] + + +class _DeepinfraRerankResponse(TypedDict, total=False): + """Body of a DeepInfra ``/rerank`` response.""" + + scores: ReadOnly[Sequence[float]] + input_tokens: ReadOnly[int] + request_id: ReadOnly[str | None] + inference_status: ReadOnly[_DeepinfraInferenceStatus] + + +class _DeepinfraRerankResponseSource(Protocol): + """The DeepInfra ``/rerank`` HTTP response, read for the body it decodes to.""" + + def json(self) -> _DeepinfraRerankResponse: ... + + +def _deepinfra_rerank_body(response: _DeepinfraRerankResponseSource) -> _DeepinfraRerankResponse: + """Decode the body of a DeepInfra ``/rerank`` response.""" + return response.json() + + class DeepinfraRerankConfig(BaseRerankConfig): """ Deepinfra Rerank - Follows the same Spec as Cohere Rerank @@ -95,7 +126,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: list[str | dict[str, Any]], + documents: list[str | dict[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: list[str] | None = None, @@ -150,7 +181,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): litellm_params: dict = {}, ) -> RerankResponse: try: - response_json: Final = raw_response.json() + response_json: Final = _deepinfra_rerank_body(raw_response) logging_obj.post_call(original_response=raw_response.text) # Extract the scores from the response diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bd2b124605c..78e6e6aaf82 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -2,7 +2,7 @@ import base64 import datetime import json import math -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -128,24 +128,35 @@ def is_gemini_image_model(model: str) -> bool: return "gemini" in base_model +def _parse_image_config_string(raw_image_config: str, model: str) -> object: + try: + return json.loads(raw_image_config) + except json.JSONDecodeError as exc: + raise litellm.UnsupportedParamsError( + model=model, + message="`imageConfig` must be valid JSON when provided as a string.", + ) from exc + + def map_openai_image_params_to_gemini( - params: dict[str, Any], + params: Mapping[str, object], model: str, supported_params: Sequence[str], - optional_params: dict[str, Any] | None = None, + optional_params: Mapping[str, object] | None = None, parse_image_config_string: bool = False, -) -> dict[str, Any]: - optional_params = optional_params or {} +) -> dict[str, object]: + already_mapped: Final[Mapping[str, object]] = optional_params or {} filtered_params: Final = {key: value for key, value in params.items() if key in supported_params} - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} - if "n" in filtered_params and "n" not in optional_params: + if "n" in filtered_params and "n" not in already_mapped: mapped_params["sampleCount"] = filtered_params["n"] - if "size" in filtered_params and "size" not in optional_params: + size_param: Final = filtered_params.get("size") + if isinstance(size_param, str) and "size" not in already_mapped: image_config: Final = map_openai_size_to_gemini_image_config( - filtered_params["size"], + size_param, model, ) if image_config is not None: @@ -156,33 +167,30 @@ def map_openai_image_params_to_gemini( if "imageSize" in image_config: mapped_params["imageSize"] = image_config["imageSize"] - image_config_param = filtered_params.get("imageConfig") - if isinstance(image_config_param, str) and parse_image_config_string: - try: - image_config_param = json.loads(image_config_param) - except json.JSONDecodeError as exc: - raise litellm.UnsupportedParamsError( - model=model, - message="`imageConfig` must be valid JSON when provided as a string.", - ) from exc + raw_image_config: Final = filtered_params.get("imageConfig") + image_config_param: Final[object] = ( + _parse_image_config_string(raw_image_config, model) + if isinstance(raw_image_config, str) and parse_image_config_string + else raw_image_config + ) if isinstance(image_config_param, dict): mapped_params["imageConfig"] = image_config_param for key, value in filtered_params.items(): - if key not in ("n", "size", "imageConfig", "tools", "web_search_options") and key not in optional_params: + if key not in ("n", "size", "imageConfig", "tools", "web_search_options") and key not in already_mapped: mapped_params[key] = value return mapped_params -def _dedupe_gemini_search_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: +def _dedupe_gemini_search_tools(tools: list[dict[str, object]]) -> list[dict[str, object]]: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) search_tool_keys: Final = VertexGeminiConfig._search_tool_keys() seen_search_keys: Final[set[str]] = set() - deduped_tools: Final[list[dict[str, Any]]] = [] + deduped_tools: Final[list[dict[str, object]]] = [] for tool in tools: if not isinstance(tool, dict): @@ -203,7 +211,7 @@ def _dedupe_gemini_search_tools(tools: list[dict[str, Any]]) -> list[dict[str, A return deduped_tools -def _has_gemini_search_tool(tools: list[Any]) -> bool: +def _has_gemini_search_tool(tools: list[object]) -> bool: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -213,9 +221,9 @@ def _has_gemini_search_tool(tools: list[Any]) -> bool: def map_gemini_image_tools_params( - non_default_params: dict[str, Any], - mapped_params: dict[str, Any], -) -> dict[str, Any]: + non_default_params: Mapping[str, object], + mapped_params: Mapping[str, object], +) -> dict[str, object]: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -239,21 +247,24 @@ def map_gemini_image_tools_params( gemini_config._drop_search_tools_mixed_with_functions(result) - if isinstance(result.get("tools"), list): - result["tools"] = _dedupe_gemini_search_tools(result["tools"]) + resolved_tools: Final = result.get("tools") + if isinstance(resolved_tools, list): + result["tools"] = _dedupe_gemini_search_tools(resolved_tools) return result def get_gemini_image_web_search_requests( - response_data: dict[str, Any], + response_data: Mapping[str, object], ) -> int | None: from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) - grounding_metadata: Final[list[dict[str, Any]]] = [] - for candidate in response_data.get("candidates", []): + raw_candidates: Final = response_data.get("candidates") + candidates: Final[list[object]] = raw_candidates if isinstance(raw_candidates, list) else [] + grounding_metadata: Final[list[dict[str, object]]] = [] + for candidate in candidates: if not isinstance(candidate, dict): continue candidate_grounding = candidate.get("groundingMetadata") @@ -267,13 +278,14 @@ def get_gemini_image_web_search_requests( def get_gemini_image_generation_config( model: str, - optional_params: dict[str, Any], -) -> dict[str, Any]: - generation_config: Final[dict[str, Any]] = {"response_modalities": ["IMAGE", "TEXT"]} + optional_params: Mapping[str, object], +) -> dict[str, object]: + generation_config: Final[dict[str, object]] = {"response_modalities": ["IMAGE", "TEXT"]} - image_config: Final[dict[str, Any]] = {} - if isinstance(optional_params.get("imageConfig"), dict): - image_config.update(optional_params["imageConfig"]) + raw_image_config: Final = optional_params.get("imageConfig") + image_config: Final[dict[str, object]] = {} + if isinstance(raw_image_config, dict): + image_config.update(raw_image_config) if not supports_gemini_image_size(model): image_config.pop("imageSize", None) @@ -398,7 +410,7 @@ class GeminiModelInfo(BaseLLMModelInfo): f"Failed to fetch models from Gemini. Status code: {response.status_code}, Response: {response.json()}" ) - models: Final = response.json()["models"] + models: Final[list[dict[str, str]]] = response.json()["models"] litellm_model_names: Final = self.process_model_name(models) return litellm_model_names @@ -473,12 +485,12 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): async def count_tokens( self, model_to_use: str, - messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index dee83407cb5..2c62e04c5a3 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -5,11 +5,13 @@ For vertex ai, check out the vertex_ai/files/handler.py file. """ import time -from typing import Any, Final, Literal +from collections.abc import Mapping +from typing import Final, Literal, TypedDict from urllib.parse import urlparse import httpx from openai.types.file_deleted import FileDeleted +from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data @@ -18,7 +20,6 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) -from litellm.types.llms.gemini import GeminiCreateFilesResponseObject from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -31,6 +32,25 @@ from litellm.types.utils import LlmProviders from ..common_utils import GeminiModelInfo +class _GeminiFileMetadata(TypedDict, total=False): + name: ReadOnly[str] + uri: ReadOnly[Required[str]] + displayName: ReadOnly[Required[str]] + mimeType: ReadOnly[str] + sizeBytes: ReadOnly[Required[str]] + createTime: ReadOnly[Required[str]] + updateTime: ReadOnly[str] + expirationTime: ReadOnly[str] + sha256Hash: ReadOnly[str] + state: ReadOnly[str] + source: ReadOnly[str] + error: ReadOnly[Mapping[str, object]] + + +class _GeminiCreateFileResponse(TypedDict): + file: ReadOnly[_GeminiFileMetadata] + + class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def __init__(self): pass @@ -41,14 +61,14 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): def validate_environment( self, - headers: dict[Any, Any], + headers: dict[str, str], model: str, messages: list[AllMessageValues], - optional_params: dict[Any, Any], - litellm_params: dict[Any, Any], + optional_params: dict[str, object], + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict[Any, Any]: + ) -> dict[str, str]: """ Validate environment and add Gemini API key to headers. Google AI Studio uses x-goog-api-key header for authentication. @@ -164,9 +184,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): Transform Gemini's file upload response into OpenAI-style FileObject """ try: - response_json: Final = raw_response.json() + response_json: Final[_GeminiCreateFileResponse] = raw_response.json() - response_object: Final = GeminiCreateFilesResponseObject(**response_json.get("file", {})) + response_object: Final = response_json["file"] # Extract file information from Gemini response @@ -262,7 +282,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ try: verbose_logger.debug("Retrieve file response: %s", raw_response.text) - response_json: Final = raw_response.json() + response_json: Final[_GeminiFileMetadata] = raw_response.json() verbose_logger.debug("Response JSON: %s", response_json) # Map Gemini state to OpenAI status gemini_state: Final = response_json.get("state", "STATE_UNSPECIFIED") diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index dcd2e4e3471..6d0f211ed7b 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -12,9 +12,10 @@ Schema versioning: litellm.use_legacy_interactions_schema = True. Remove flag after June 8, 2026. """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -41,6 +42,53 @@ else: LiteLLMLoggingObj = Any +_JsonObject: TypeAlias = dict[str, object] + + +class _InteractionPayload(TypedDict, total=False): + """JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields.""" + + id: ReadOnly[str | None] + object: ReadOnly[str | None] + model: ReadOnly[str | None] + agent: ReadOnly[str | None] + status: ReadOnly[str | None] + created: ReadOnly[str | None] + updated: ReadOnly[str | None] + outputs: ReadOnly[list[_JsonObject] | None] + steps: ReadOnly[list[_JsonObject] | None] + usage: ReadOnly[_JsonObject | None] + + +class _CancelPayload(TypedDict, total=False): + """JSON body of an Interactions API cancel response.""" + + id: ReadOnly[str | None] + status: ReadOnly[str | None] + + +class _InteractionPayloadSource(Protocol): + """An Interactions API HTTP response, read for the interaction body it decodes to.""" + + def json(self) -> _InteractionPayload: ... + + +class _CancelPayloadSource(Protocol): + """An Interactions API cancel HTTP response, read for the body it decodes to.""" + + def json(self) -> _CancelPayload: ... + + +def _interaction_body(response: _InteractionPayloadSource) -> _InteractionPayload: + """Decode the body of an Interactions API interaction response.""" + return response.json() + + +def _cancel_body(response: _CancelPayloadSource) -> _CancelPayload: + """Decode the body of an Interactions API cancel response.""" + return response.json() + + class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ Configuration for Google AI Studio Interactions API. @@ -143,7 +191,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ use_legacy: Final[bool] = litellm.use_legacy_interactions_schema - request_body: Final[dict[str, Any]] = {} + request_body: Final[dict[str, object]] = {} # Model or Agent (one required) if model: @@ -189,7 +237,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): and (not isinstance(response_format, dict) or "mime_type" not in response_format) ): # Wrap the legacy schema into the new polymorphic format. - new_rf: Final[dict[str, Any]] = { + new_rf: Final[dict[str, object]] = { "type": "text", "mime_type": response_mime_type, } @@ -215,7 +263,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if image_config is not None: # Move image_config to response_format with type=image. - image_rf: Final[dict[str, Any]] = {"type": "image", **image_config} + image_rf: Final[_JsonObject] = {"type": "image", **image_config} existing_rf: Final = request_body.get("response_format") if existing_rf is None: request_body["response_format"] = image_rf @@ -239,7 +287,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): original_response=raw_response.text, additional_args={"complete_input_dict": {}}, ) - raw_json: Final = raw_response.json() + raw_json: Final = _interaction_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, @@ -290,7 +338,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> InteractionsAPIResponse: try: - raw_json: Final = raw_response.json() + raw_json: Final = _interaction_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, @@ -355,7 +403,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> CancelInteractionResult: try: - raw_json: Final = raw_response.json() + raw_json: Final = _cancel_body(raw_response) except Exception: raise GeminiError( message=raw_response.text, diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 367619db37d..c92af7de145 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -7,6 +7,8 @@ from collections import OrderedDict from collections.abc import Mapping, Sequence from typing import Any, Final, cast +from typing_extensions import ReadOnly, Required, TypedDict + import litellm from litellm import verbose_logger from litellm._uuid import uuid @@ -96,6 +98,23 @@ def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None: return VertexGeminiConfig()._map_audio_params({"voice": voice}) +class _GeminiLiveSetupEnvelope(TypedDict, total=False): + setup: ReadOnly[BidiGenerateContentSetup] + + +class _OpenAIRealtimeClientEvent(TypedDict, total=False): + type: ReadOnly[str] + audio: ReadOnly[Required[str]] + session: ReadOnly[dict[str, object]] + item: ReadOnly[dict[str, object]] + + +def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup: + envelope: Final[_GeminiLiveSetupEnvelope] = json.loads(session_configuration_request) + empty_setup: Final[BidiGenerateContentSetup] = {} + return envelope.get("setup", empty_setup) + + # Google bills Live transcription at an estimated 25 audio tokens/sec of input and # 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing). GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25 @@ -130,7 +149,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return True @staticmethod - def _usage_detail_alias(details: Any, defaults: dict[str, int]) -> dict[str, Any]: + def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: if not isinstance(details, dict): return dict(defaults) return { @@ -139,7 +158,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): } @staticmethod - def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, Any]: + def _add_pipecat_usage_detail_aliases(usage_dict: dict[str, Any]) -> dict[str, object]: usage_dict.setdefault( "input_token_details", GeminiRealtimeConfig._usage_detail_alias( @@ -222,8 +241,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if not session_configuration_request: return False try: - setup: Final = json.loads(session_configuration_request).get("setup", {}) - automatic_detection: Final = setup.get("realtimeInputConfig", {}).get("automaticActivityDetection", {}) + setup: Final = _parse_setup(session_configuration_request) + automatic_detection: Final[object] = setup.get("realtimeInputConfig", {}).get( + "automaticActivityDetection", {} + ) return isinstance(automatic_detection, dict) and automatic_detection.get("disabled") is True except (json.JSONDecodeError, TypeError, AttributeError): return False @@ -406,7 +427,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return "TEXT" if GeminiRealtimeConfig._is_text_only_live_model(model) else "AUDIO" @staticmethod - def _coerce_response_modalities(model: str, modalities: Sequence[Any]) -> tuple[str, ...]: + def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]: """Swap responseModalities a Live model cannot produce: TEXT to AUDIO for audio-only models, AUDIO to TEXT for text-only ones (e.g. transcribe-live).""" normalized: Final = tuple( @@ -431,7 +452,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): def _handle_session_update( self, - json_message: dict, + json_message: _OpenAIRealtimeClientEvent, model: str, session_configuration_request: str | None, ) -> list[str]: @@ -445,7 +466,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): with a 1007, tearing the session down). To carry tools/instructions, send them on the first session.update before any conversation content. """ - session_payload = json_message.get("session") or {} + empty_session: Final[dict[str, object]] = {} + session_payload = json_message.get("session") or empty_session # Normalize GA-remapped fields (``output_modalities``, # nested ``audio.input.transcription``, # ``audio.input.turn_detection``) back to their flat beta keys so @@ -486,14 +508,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Gemini Realtime: Ignoring session.update (setup already sent)") return [] - def _handle_conversation_item(self, json_message: dict) -> list[str]: + def _handle_conversation_item(self, json_message: _OpenAIRealtimeClientEvent) -> list[str]: """ Handle conversation.item.create for user text or function call output. Converts OpenAI format to Gemini's clientContent (for user text) or toolResponse (for function outputs). """ - item: Final = json_message.get("item", {}) + empty_item: Final[dict[str, object]] = {} + item: Final = json_message.get("item", empty_item) item_type: Final = item.get("type") if item_type == "function_call_output": @@ -524,7 +547,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): call_id, ) - function_response: Final[dict[str, Any]] = {"response": output_dict} + function_response: Final[dict[str, object]] = {"response": output_dict} if self._include_function_response_id() and call_id: function_response["id"] = call_id if function_name: @@ -559,7 +582,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) -> list[str]: realtime_input_dict: BidiGenerateContentRealtimeInput = {} try: - json_message: Final = json.loads(message) + json_message: Final[_OpenAIRealtimeClientEvent] = json.loads(message) except json.JSONDecodeError: if isinstance(message, bytes): message_str = message.decode("utf-8", errors="replace") @@ -610,9 +633,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request: str | None = None, ) -> OpenAIRealtimeStreamSessionEvents: if session_configuration_request: - session_configuration_request_dict: BidiGenerateContentSetup = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request) else: session_configuration_request_dict = {} @@ -663,7 +684,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_configuration_request_dict: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_configuration_request_dict = json.loads(session_configuration_request).get("setup", {}) + session_configuration_request_dict = _parse_setup(session_configuration_request) except json.JSONDecodeError: session_configuration_request_dict = {} generation_config: Final = session_configuration_request_dict.get("generationConfig", {}) @@ -931,9 +952,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): return events @staticmethod - def get_nested_value(obj: dict, path: str) -> Any: + def get_nested_value(obj: dict, path: str) -> object | None: keys: Final = path.split(".") - current = obj + current: object = obj for key in keys: if isinstance(current, dict) and key in current: current = current[key] @@ -1011,9 +1032,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): current_response_id = f"resp_{uuid.uuid4()}" if session_configuration_request: - session_configuration_request_dict: BidiGenerateContentSetup = json.loads( - session_configuration_request - ).get("setup", {}) + session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request) else: session_configuration_request_dict = {} @@ -1337,7 +1356,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): session_setup: BidiGenerateContentSetup = {} if session_configuration_request is not None: try: - session_setup = json.loads(session_configuration_request).get("setup", {}) + session_setup = _parse_setup(session_configuration_request) except (json.JSONDecodeError, TypeError): session_setup = {} tool_call_generation_config = session_setup.get("generationConfig", {}) or {} diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 6a1fc144c42..ff4c675b02f 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -1,4 +1,5 @@ import base64 +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -54,8 +55,13 @@ def _convert_image_to_gemini_format(image_file) -> dict[str, str]: return {"bytesBase64Encoded": base64_encoded, "mimeType": mime_type} +def _json_payload(raw_response: httpx.Response) -> object: + """Read an HTTP response body as an opaque JSON payload.""" + return raw_response.json() + + def _usage_video_resolution_from_parameters( - parameters: dict[str, Any], + parameters: Mapping[str, object], ) -> str | None: """Normalize Veo ``parameters.resolution`` for usage and cost tracking.""" res: Final = parameters.get("resolution") @@ -97,7 +103,7 @@ class GeminiVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Map OpenAI-style parameters to Veo format. @@ -111,7 +117,7 @@ class GeminiVideoConfig(BaseVideoConfig): All other params are passed through as-is to support Gemini-specific parameters. """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Get supported OpenAI params (exclude "model" and "prompt" which are handled separately) supported_openai_params: Final = self.get_supported_openai_params(model) @@ -312,11 +318,11 @@ class GeminiVideoConfig(BaseVideoConfig): - status: "processing" - usage: includes duration_seconds and optional video_resolution for cost calculation """ - response_data: Final = raw_response.json() + response_data: Final = _json_payload(raw_response) # Parse response using Pydantic model for type safety try: - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) except Exception as e: raise ValueError(f"Failed to parse operation response: {e}") @@ -336,7 +342,7 @@ class GeminiVideoConfig(BaseVideoConfig): model=model, ) - usage_data: Final[dict[str, Any]] = {} + usage_data: Final[dict[str, float | str]] = {} if request_data: parameters: Final = request_data.get("parameters", {}) duration: Final = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS @@ -367,7 +373,7 @@ class GeminiVideoConfig(BaseVideoConfig): """ operation_name: Final = extract_original_video_id(video_id) url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return url, params @@ -403,9 +409,9 @@ class GeminiVideoConfig(BaseVideoConfig): } } """ - response_data: Final = raw_response.json() + response_data: Final = _json_payload(raw_response) # Parse response using Pydantic model for type safety - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) operation_name: Final = operation_response.name is_done: Final = operation_response.done @@ -443,9 +449,9 @@ class GeminiVideoConfig(BaseVideoConfig): client: Final = litellm.module_level_client status_response: Final = client.get(url=status_url, headers=headers) status_response.raise_for_status() - response_data: Final = status_response.json() + response_data: Final = _json_payload(status_response) - operation_response: Final = GeminiLongRunningOperationResponse(**response_data) + operation_response: Final = GeminiLongRunningOperationResponse.model_validate(response_data) if not operation_response.done: raise ValueError( @@ -458,7 +464,7 @@ class GeminiVideoConfig(BaseVideoConfig): generated_samples: Final = operation_response.response.generateVideoResponse.generatedSamples download_url: Final = generated_samples[0].video.uri - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} return download_url, params @@ -480,7 +486,7 @@ class GeminiVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Video remix is not supported by Veo API. @@ -506,7 +512,7 @@ class GeminiVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Video list is not supported by Veo API. @@ -547,7 +553,7 @@ class GeminiVideoConfig(BaseVideoConfig): """Video delete is not supported.""" raise NotImplementedError("Video delete is not supported by Google Veo.") - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for Gemini") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index 2c58e16fc23..50474aabdeb 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -4,13 +4,12 @@ VLLM is a superset of OpenAI's `embedding` endpoint. ## `encoding_format` -For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request: +For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request. `hosted_vllm/...` models use a separate handler that never adds the field on its own, so this resolution applies to the `openai/...`-style routes only: 1. Explicit value on the embedding call (`encoding_format=...`). 2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). 3. Environment variable `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT` (e.g. in `.env` or container env). -4. Default **`float`**. -That avoids forwarding `encoding_format=None` to the provider/SDK where some servers behave poorly. +If none of those is set, or the winning value is the literal string `none`, the field is omitted from the upstream request entirely (LiteLLM also bypasses the OpenAI SDK's own base64 default), so OpenAI-compatible servers that reject `encoding_format` keep working. -To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). \ No newline at end of file +To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index d3db3530109..f6fe7f2fa10 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -1,8 +1,9 @@ import json import os import time +from collections.abc import Sequence from copy import deepcopy -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx @@ -24,6 +25,8 @@ from litellm.utils import token_counter from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -31,6 +34,12 @@ else: LoggingClass = Any +class _TokenEncoding(Protocol): + """Tokenizer handle the caller passes in; only `encode` is used, to count completion tokens.""" + + def encode(self, text: str, /) -> Sequence[object]: ... + + tgi_models_cache = None conv_models_cache = None @@ -369,7 +378,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): model_response: ModelResponse, task: hf_tasks | None, optional_params: dict, - encoding: Any, + encoding: "_TokenEncoding | None", messages: list[AllMessageValues], model: str, ): @@ -439,9 +448,10 @@ class HuggingFaceEmbeddingConfig(BaseConfig): if output_text is not None and len(output_text) > 0: completion_tokens = 0 try: - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) - ) ##[TODO] use the llama2 tokenizer here + if encoding is not None: + completion_tokens = len( + encoding.encode(model_response["choices"][0]["message"].get("content", "")) + ) ##[TODO] use the llama2 tokenizer here except Exception: # this should remain non blocking we should not block a response returning if calculating usage fails pass @@ -469,7 +479,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index f51213ca1c3..fbc287589b3 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -13,55 +13,109 @@ Generated files are returned directly in the response - no separate storage need import base64 import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum -from typing import Any, Final, Protocol +from typing import Any, Final, Protocol, TypedDict -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ReadOnly from litellm._logging import verbose_logger -class _ToolCallFunction(Protocol): - """Function payload of an assistant tool call.""" - - name: str | None - arguments: str +class _ToolParameterSchema(TypedDict, total=False): + type: ReadOnly[str] + description: ReadOnly[str] -class _ToolCall(Protocol): - """Tool call requested by the assistant on a chat completion choice.""" - - id: str - function: _ToolCallFunction +class _ToolArgumentSchema(TypedDict, total=False): + type: ReadOnly[str] + properties: ReadOnly[Mapping[str, _ToolParameterSchema]] + required: ReadOnly[Sequence[str]] -class _AssistantMessage(Protocol): - """Assistant message carried by a chat completion choice.""" - - content: str | None - tool_calls: Sequence[_ToolCall] | None +class _OpenAIToolFunction(TypedDict, total=False): + name: ReadOnly[str] + description: ReadOnly[str] + parameters: ReadOnly[_ToolArgumentSchema] -class _CompletionChoice(Protocol): - """Single choice of a chat completion response.""" - - finish_reason: str - message: _AssistantMessage +class _OpenAIToolSpec(TypedDict, total=False): + type: ReadOnly[str] + function: ReadOnly[_OpenAIToolFunction] -class _SandboxFile(TypedDict): - """File generated inside the sandbox during a code execution run.""" +class _AnthropicToolSpec(TypedDict, total=False): + name: ReadOnly[str] + description: ReadOnly[str] + input_schema: ReadOnly[_ToolArgumentSchema] + +class _CodeExecutionArguments(TypedDict, total=False): + code: ReadOnly[str] + + +class _GeneratedFile(TypedDict, total=False): + name: ReadOnly[str] + mime_type: ReadOnly[str] + content_base64: ReadOnly[str] + size: ReadOnly[int] + + +class _SandboxGeneratedFile(TypedDict): name: ReadOnly[str] mime_type: ReadOnly[str] content_base64: ReadOnly[str] -class _CodeExecutionArguments(TypedDict): - """Arguments the model passes to the `litellm_code_execution` tool.""" +class _SandboxExecutionResult(TypedDict): + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[_SandboxGeneratedFile]] - code: NotRequired[ReadOnly[str]] + +class _ExecutionResult(TypedDict, total=False): + iteration: ReadOnly[int] + success: ReadOnly[bool] + output: ReadOnly[str] + error: ReadOnly[str] + files: ReadOnly[Sequence[str]] + + +class _ToolCallFunction(Protocol): + name: str + arguments: str + + +class _ToolCall(Protocol): + id: str + function: _ToolCallFunction + + +class _AssistantMessage(Protocol): + content: str | None + tool_calls: Sequence[_ToolCall] | None + + +class _ResponseChoice(Protocol): + message: _AssistantMessage + finish_reason: str | None + + +class _CompletionResponse(Protocol): + choices: Sequence[_ResponseChoice] + + +class _CodeExecutionOutcome(TypedDict, total=False): + response: ReadOnly[_CompletionResponse | None] + files: ReadOnly[Sequence[_GeneratedFile]] + execution_results: ReadOnly[Sequence[_ExecutionResult]] + messages: ReadOnly[Sequence[dict[str, object]]] + max_iterations_reached: ReadOnly[bool] + + +def _parse_code_execution_arguments(serialized_arguments: str) -> _CodeExecutionArguments: + return json.loads(serialized_arguments) class LiteLLMInternalTools(str, Enum): @@ -75,7 +129,7 @@ class LiteLLMInternalTools(str, Enum): CODE_EXECUTION = "litellm_code_execution" -def get_litellm_code_execution_tool() -> dict[str, object]: +def get_litellm_code_execution_tool() -> _OpenAIToolSpec: """ Returns the litellm_code_execution tool definition in OpenAI format. @@ -96,7 +150,7 @@ def get_litellm_code_execution_tool() -> dict[str, object]: } -def get_litellm_code_execution_tool_anthropic() -> dict[str, object]: +def get_litellm_code_execution_tool_anthropic() -> _AnthropicToolSpec: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. @@ -143,12 +197,12 @@ class CodeExecutionHandler: async def execute_with_code_execution( self, model: str, - messages: list[dict], - tools: list[dict], + messages: list[dict[str, object]], + tools: list[_OpenAIToolSpec], skill_files: dict[str, bytes], skill_id: str | None = None, **kwargs, - ) -> dict[str, object]: + ) -> _CodeExecutionOutcome: """ Execute an LLM call with automatic code execution handling. @@ -179,8 +233,8 @@ class CodeExecutionHandler: ) current_messages: Final = list(messages) - generated_files: Final[list[dict[str, object]]] = [] # Files returned directly - execution_results: Final[list[dict[str, object]]] = [] + generated_files: Final[list[_GeneratedFile]] = [] # Files returned directly + execution_results: Final[list[_ExecutionResult]] = [] executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error @@ -196,9 +250,9 @@ class CodeExecutionHandler: **kwargs, ) - choice: _CompletionChoice = response.choices[0] + choice: _ResponseChoice = response.choices[0] assistant_message = choice.message - stop_reason: str = choice.finish_reason + stop_reason = choice.finish_reason # Build assistant message for conversation history assistant_msg_dict: dict[str, object] = { @@ -236,19 +290,19 @@ class CodeExecutionHandler: if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: - args: _CodeExecutionArguments = json.loads(tool_call.function.arguments) - code: str = args.get("code", "") + args = _parse_code_execution_arguments(tool_call.function.arguments) + code = args.get("code", "") verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code)) - exec_result = executor.execute( + exec_result: _SandboxExecutionResult = executor.execute( code=code, skill_files=skill_files, ) verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result) - sandbox_files: Sequence[_SandboxFile] = exec_result["files"] + sandbox_files: Sequence[_SandboxGeneratedFile] = exec_result["files"] execution_results.append( { @@ -326,7 +380,7 @@ class CodeExecutionHandler: } -def has_code_execution_tool(tools: list[dict] | None) -> bool: +def has_code_execution_tool(tools: list[_OpenAIToolSpec] | None) -> bool: """Check if litellm_code_execution tool is in the tools list.""" if not tools: return False @@ -337,7 +391,7 @@ def has_code_execution_tool(tools: list[dict] | None) -> bool: return False -def add_code_execution_tool(tools: list[dict] | None) -> list[dict]: +def add_code_execution_tool(tools: list[_OpenAIToolSpec] | None) -> list[_OpenAIToolSpec]: """Add litellm_code_execution tool if not already present.""" tools = tools or [] if not has_code_execution_tool(tools): diff --git a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py index 008a5a5780f..046b4e29a0a 100644 --- a/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py +++ b/litellm/llms/nvidia_riva/audio_transcription/audio_utils.py @@ -16,7 +16,7 @@ import io import os import tempfile from dataclasses import dataclass -from typing import Any, Final, cast +from typing import Final, Protocol, cast from litellm.llms.nvidia_riva.audio_transcription.transformation import ( RIVA_TARGET_NUM_CHANNELS, @@ -24,10 +24,30 @@ from litellm.llms.nvidia_riva.audio_transcription.transformation import ( ) from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException -# Keep this as Any: the module intentionally avoids importing numpy at module -# import time (optional dependency), and project-wide mypy config evaluates this -# file in contexts where conditional type aliases can degrade to "FloatArray?". -FloatArray = Any + +class FloatArray(Protocol): + """Structural view of the ``numpy.ndarray`` surface this module relies on.""" + + @property + def ndim(self) -> int: ... + + @property + def shape(self) -> tuple[int, ...]: ... + + @property + def size(self) -> int: ... + + def mean(self, axis: int) -> "FloatArray": ... + + def ravel(self) -> "FloatArray": ... + + def astype(self, dtype: object) -> "FloatArray": ... + + def tobytes(self) -> bytes: ... + + def __getitem__(self, key: object) -> "FloatArray": ... + + def __mul__(self, other: float) -> "FloatArray": ... _INSTALL_HINT = "Install Riva STT extras to enable automatic audio resampling: `pip install 'litellm[stt-nvidia-riva]'`" diff --git a/litellm/llms/oci/common_utils.py b/litellm/llms/oci/common_utils.py index 5c3962bc05d..3f703564b5a 100644 --- a/litellm/llms/oci/common_utils.py +++ b/litellm/llms/oci/common_utils.py @@ -5,10 +5,11 @@ import os import re from dataclasses import dataclass from email.utils import formatdate -from typing import Any, Final, Protocol +from typing import Final, Protocol from urllib.parse import urlparse import httpx +from pydantic import JsonValue from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -64,7 +65,7 @@ class OCISignerProtocol(Protocol): See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html """ - def do_request_sign(self, request: Any, *, enforce_content_headers: bool = False) -> None: + def do_request_sign(self, request: "OCIRequestWrapper", *, enforce_content_headers: bool = False) -> None: pass @@ -113,7 +114,7 @@ def build_signature_string(method: str, path: str, headers: dict, signed_headers return "\n".join(lines) -def load_private_key_from_str(key_str: str) -> Any: +def load_private_key_from_str(key_str: str) -> "rsa.RSAPrivateKey": _require_cryptography() key: Final = serialization.load_pem_private_key( key_str.encode("utf-8"), @@ -124,7 +125,7 @@ def load_private_key_from_str(key_str: str) -> Any: return key -def load_private_key_from_file(file_path: str) -> Any: +def load_private_key_from_file(file_path: str) -> "rsa.RSAPrivateKey": """Loads a private key from a file path.""" try: with open(file_path, "r", encoding="utf-8") as f: @@ -421,16 +422,17 @@ OCI_JSON_TO_PYTHON_TYPES: Final[dict[str, str]] = { } -def resolve_oci_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: +def resolve_oci_schema_refs(schema: JsonValue) -> JsonValue: """Inline all ``$ref``/``$defs`` references — OCI does not support JSON Schema ``$ref``.""" - defs: Final = schema.get("$defs", {}) - resolving_stack: Final[set] = set() + raw_defs: Final = schema.get("$defs") if isinstance(schema, dict) else None + defs: Final[dict[str, JsonValue]] = raw_defs if isinstance(raw_defs, dict) else {} + resolving_stack: Final[set[str]] = set() - def _resolve(obj: Any) -> Any: + def _resolve(obj: JsonValue) -> JsonValue: if isinstance(obj, dict): - if "$ref" in obj: - ref: Final = obj["$ref"] - if ref.startswith("#/$defs/"): + ref: Final = obj.get("$ref") + if ref is not None: + if isinstance(ref, str) and ref.startswith("#/$defs/"): key: Final = ref.split("/")[-1] if key in resolving_stack: return {"type": "object"} # break cycles @@ -451,7 +453,7 @@ def resolve_oci_schema_refs(schema: dict[str, Any]) -> dict[str, Any]: return resolved -def resolve_oci_schema_anyof(obj: Any) -> Any: +def resolve_oci_schema_anyof(obj: JsonValue) -> JsonValue: """Resolve Pydantic v2 ``Optional[T]`` → ``anyOf`` patterns. Pydantic v2 emits ``{"anyOf": [{"type": "T"}, {"type": "null"}]}`` for @@ -459,10 +461,13 @@ def resolve_oci_schema_anyof(obj: Any) -> Any: first non-null branch and merge top-level metadata into it. """ if isinstance(obj, dict): - if "anyOf" in obj and "type" not in obj: - non_null: Final = [t for t in obj["anyOf"] if not (isinstance(t, dict) and t.get("type") == "null")] + raw_any_of: Final = obj.get("anyOf") + if raw_any_of is not None and "type" not in obj: + branches: Final = raw_any_of if isinstance(raw_any_of, list) else [] + non_null: Final = [t for t in branches if not (isinstance(t, dict) and t.get("type") == "null")] if non_null: - resolved: Final = {**obj, **non_null[0]} + first: Final = non_null[0] + resolved: Final[dict[str, JsonValue]] = {**obj, **first} if isinstance(first, dict) else {**obj} resolved.pop("anyOf", None) return resolve_oci_schema_anyof(resolved) return {k: resolve_oci_schema_anyof(v) for k, v in obj.items()} @@ -471,7 +476,7 @@ def resolve_oci_schema_anyof(obj: Any) -> Any: return obj -def sanitize_oci_schema(schema: Any) -> Any: +def sanitize_oci_schema(schema: JsonValue) -> JsonValue: """Recursively remove OCI-incompatible fields from a JSON schema. Strips ``title`` keys, removes ``None``-valued ``default`` entries, @@ -483,7 +488,7 @@ def sanitize_oci_schema(schema: Any) -> Any: if not isinstance(schema, dict): return schema - sanitized: Final[dict[str, Any]] = {} + sanitized: Final[dict[str, JsonValue]] = {} for key, value in schema.items(): if key == "title": continue @@ -513,7 +518,7 @@ def sanitize_oci_schema(schema: Any) -> Any: return sanitized -def enrich_cohere_param_description(description: str, param_schema: dict[str, Any]) -> str: +def enrich_cohere_param_description(description: str, param_schema: dict[str, JsonValue]) -> str: """Embed schema constraints into a Cohere parameter description. ``CohereParameterDefinition`` only has ``type``, ``description``, and diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 5894658e5d2..9afc6331d96 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -4,7 +4,8 @@ Support for gpt model family import json import os -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload from urllib.parse import urlparse @@ -21,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, get_tool_call_names, hoist_images_from_tool_messages, + tool_with_flattened_parameters, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -65,6 +67,9 @@ else: LiteLLMLoggingObj = Any +_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) + + class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Reference: https://platform.openai.com/docs/api-reference/chat/create @@ -170,16 +175,20 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if model != "gpt-3.5-turbo-16k" and model != "gpt-4": # gpt-4 does not support 'response_format' model_specific_params.append("response_format") - # Normalize model name for responses API (e.g., "responses/gpt-4.1" -> "gpt-4.1") - model_for_check: Final = model.split("responses/", 1)[1] if "responses/" in model else model - if ( - model_for_check in litellm.open_ai_chat_completion_models - ) or model_for_check in litellm.open_ai_text_completion_models: + if OpenAIGPTConfig.is_openai_catalog_model(model): model_specific_params.append( "user" ) # user is not a param supported by all openai-compatible endpoints - e.g. azure ai return base_params + model_specific_params + @staticmethod + def is_openai_catalog_model(model: str) -> bool: + model_for_check: Final = model.split("responses/", 1)[1] if "responses/" in model else model + return ( + model_for_check in litellm.open_ai_chat_completion_models + or model_for_check in litellm.open_ai_text_completion_models + ) + def _map_openai_params( self, non_default_params: dict, @@ -321,7 +330,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] - ) -> Coroutine[Any, Any, list[AllMessageValues]]: + ) -> Coroutine[object, object, list[AllMessageValues]]: ... @overload @@ -337,7 +346,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: bool = False - ) -> list[AllMessageValues] | Coroutine[Any, Any, list[AllMessageValues]]: + ) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: """OpenAI no longer supports image_url as a string, so we need to convert it to a dict""" stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages) hoisted_messages: Final = hoist_images_from_tool_messages(stripped_messages) @@ -393,6 +402,21 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) return messages, tools + def _targets_openai_hosted_endpoint( + self, + custom_llm_provider: str | None, + api_base: str | None, + ) -> bool: + if custom_llm_provider != "openai": + return False + resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") + if not resolved_api_base: + return True + hostname: Final = urlparse(resolved_api_base).hostname + if hostname is None: + return True + return hostname == "openai.com" or hostname.endswith(".openai.com") + def _should_preserve_cache_control_for_endpoint( self, custom_llm_provider: str | None, @@ -404,15 +428,34 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): api_base. Those can understand cache_control, so it must survive there. Real OpenAI cannot, so it is still stripped for an openai.com host. """ - if custom_llm_provider != "openai": - return False - resolved_api_base = api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE") - if not resolved_api_base: - return False - hostname: Final = urlparse(resolved_api_base).hostname - if hostname is None: - return False - return hostname != "openai.com" and not hostname.endswith(".openai.com") + return custom_llm_provider == "openai" and not self._targets_openai_hosted_endpoint( + custom_llm_provider, api_base + ) + + def _flattened_tools_update_for_openai( + self, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> Mapping[str, object]: + """ + OpenAI's chat completions validator rejects tool `parameters` carrying + 'oneOf'/'anyOf'/'allOf'/'enum'/'const'/'not' at the top level for every + model family, unlike the Responses API, where GPT-5+ accepts them. + """ + tools: Final = optional_params.get("tools") + if not isinstance(tools, list): + return _NO_TOOLS_UPDATE + provider: Final = litellm_params.get("custom_llm_provider") + raw_api_base: Final = litellm_params.get("api_base") + if not self._targets_openai_hosted_endpoint( + provider if isinstance(provider, str) else None, + raw_api_base if isinstance(raw_api_base, str) else None, + ): + return _NO_TOOLS_UPDATE + flattened: Final = [ # mutable-ok: request tools are a JSON list + tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + ] + return MappingProxyType({"tools": flattened}) def transform_request( self, @@ -439,11 +482,14 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): optional_params["tools"] = tools optional_params.pop("max_retries", None) + if not optional_params.get("tools") and not optional_params.get("functions"): + optional_params.pop("tool_choice", None) return { "model": model, "messages": messages, **optional_params, + **self._flattened_tools_update_for_openai(optional_params, litellm_params), } async def async_transform_request( @@ -469,10 +515,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if tools is not None and len(tools) > 0: optional_params["tools"] = tools if self.__class__._is_base_class: + if not optional_params.get("tools") and not optional_params.get("functions"): + optional_params.pop("tool_choice", None) return { "model": model, "messages": transformed_messages, **optional_params, + **self._flattened_tools_update_for_openai(optional_params, litellm_params), } else: ## allow for any object specific behaviour to be handled @@ -493,8 +542,12 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): return None tool_call_names: Final = get_tool_call_names(optional_params.get("tools", [])) try: - json_content: Final = json.loads(content) - if json_content.get("type") == "function" and json_content.get("name") in tool_call_names: + json_content: Final[object] = json.loads(content) + if ( + isinstance(json_content, dict) + and json_content.get("type") == "function" + and json_content.get("name") in tool_call_names + ): return ChatCompletionMessageToolCall( function=Function( name=json_content.get("name"), @@ -618,7 +671,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ## RESPONSE OBJECT try: - completion_response: Final = raw_response.json() + completion_response: Final[dict[str, object]] = raw_response.json() except Exception as e: response_headers: Final = getattr(raw_response, "headers", None) raise OpenAIError( @@ -755,6 +808,14 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) +class OpenAIUnknownModelConfig(OpenAIGPTConfig): + """A model the openai provider does not recognize is typically a LiteLLM proxy alias, so + forward reasoning_effort and let the server decide whether it is supported.""" + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: inherited contract + return super().get_supported_openai_params(model) + ["reasoning_effort"] # mutable-ok: inherited contract + + class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): def _map_reasoning_to_reasoning_content(self, choices: list) -> list: """ diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index de15fefe943..96a5ed663fc 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,9 +14,14 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +import json +import time +import uuid from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing_extensions import NotRequired, ReadOnly, TypedDict + import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( @@ -24,6 +29,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_chat_stream_usage, effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, @@ -32,6 +38,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( openai_tool_name, role_out_of_guardrail_scope, scoped_structured_message_indices, + stream_item_field, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -49,8 +56,12 @@ from litellm.types.utils import ( if TYPE_CHECKING: from fastapi import HTTPException - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth class OpenAIChatCompletionsHandler(BaseTranslation): @@ -80,7 +91,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): data: dict, guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - ) -> Any: + ) -> dict: """ Process input messages by applying guardrails to text content. """ @@ -329,9 +340,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): response: "ModelResponse", guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> Any: + ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -436,7 +447,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, ) -> list["ModelResponseStream"]: @@ -486,7 +497,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", - user_api_key_dict: Any | None, + user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can @@ -589,8 +600,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def build_stream_error_items( self, exc: "HTTPException", - responses_so_far: Sequence[Any] | None = None, - ) -> Sequence[Any] | None: + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes] | None: import json from litellm.proxy.common_request_processing import sse_error_payload @@ -630,7 +641,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", - user_api_key_dict: Any | None, + user_api_key_dict: "UserAPIKeyAuth | None", request_data: dict | None, sink: StreamTransformSink, ) -> None: @@ -794,7 +805,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Determine content source and tool calls based on choice type content = None - tool_calls: list[Any] | None = None + tool_calls: Sequence[object] | None = None if isinstance(choice, litellm.Choices): content = choice.message.content tool_calls = choice.message.tool_calls @@ -1004,3 +1015,129 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else: # Subsequent chunks - clear the text content_item["text"] = "" + + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + """ + True once any relayed chunk carries a non-null ``finish_reason``. + + The unified guardrail's ``end_of_stream_only`` streaming path probes + this via ``hasattr`` to withhold the terminal chunks until + end-of-stream moderation runs, so a block can replace the finish + instead of trailing after a ``finish_reason`` the client already saw. + """ + return any( + stream_item_field(choice, "finish_reason") is not None + for item in responses_so_far + for choice in _stream_chunk_choices(item) + ) + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + """ + Build OpenAI chat-completions SSE chunks that deliver the guardrail + block message and terminate the stream cleanly, mirroring the + non-streaming block response: ``finish_reason`` ``content_filter`` plus + the real usage the upstream call consumed. + + - ``stream_started`` False (buffered / pre-stream): nothing has been + sent, so open a standalone completion with a ``role`` delta. + - ``stream_started`` True (sampling / mid-stream): chunks already + reached the client, so continue the in-progress completion (reuse its + id/created/model, content-only delta). + + The proxy's data generator appends ``data: [DONE]`` itself. + """ + chunk_id, created, model = _blocked_stream_identity(exc, responses_so_far or ()) + prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response) + continuation_delta: Final[_BlockedChunkDelta] = {"content": exc.message} + standalone_delta: Final[_BlockedChunkDelta] = {"role": "assistant", "content": exc.message} + message_chunk: Final[_BlockedChunk] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": ( + { + "index": 0, + "delta": continuation_delta if stream_started else standalone_delta, + "finish_reason": None, + }, + ), + } + final_chunk: Final[_BlockedChunk] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": ({"index": 0, "delta": {}, "finish_reason": "content_filter"},), + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + return _chat_sse_chunk(message_chunk), _chat_sse_chunk(final_chunk) + + +class _BlockedChunkDelta(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[str] + + +class _BlockedChunkChoice(TypedDict): + index: ReadOnly[int] + delta: ReadOnly[_BlockedChunkDelta] + finish_reason: ReadOnly[str | None] + + +class _BlockedChunkUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +class _BlockedChunk(TypedDict): + id: ReadOnly[str] + object: ReadOnly[str] + created: ReadOnly[int] + model: ReadOnly[str] + choices: ReadOnly[tuple[_BlockedChunkChoice, ...]] + usage: NotRequired[ReadOnly[_BlockedChunkUsage]] + + +def _chat_sse_chunk(payload: _BlockedChunk) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _stream_chunk_choices(item: object) -> Sequence[object]: + choices: Final = stream_item_field(item, "choices") + if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)): + return choices + return () + + +def _blocked_stream_identity( + exc: "ModifyResponseException", responses_so_far: Sequence[object] +) -> tuple[str, int, str]: + identified: Final = next( + ( + (chunk_id, item) + for item in responses_so_far + if isinstance(chunk_id := stream_item_field(item, "id"), str) and chunk_id + ), + None, + ) + if identified is None: + return f"chatcmpl-{uuid.uuid4()}", int(time.time()), exc.model + chunk_id, source = identified + created: Final = stream_item_field(source, "created") + model: Final = stream_item_field(source, "model") + return ( + chunk_id, + created if isinstance(created, int) else int(time.time()), + model if isinstance(model, str) and model else exc.model, + ) diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 44bd401c115..1a5211d5ff5 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -155,10 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container creation response.""" - response_data: Final[OpenAIContainerPayload] = raw_response.json() - - # Transform the response data - container_obj: Final = ContainerObject(**response_data) + container_obj: Final = ContainerObject.model_validate(raw_response.json()) # Add cost for container creation (OpenAI containers are code interpreter sessions) # https://platform.openai.com/docs/pricing @@ -215,10 +212,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: """Transform the OpenAI container list response.""" - response_data: Final[OpenAIContainerListPayload] = raw_response.json() - - # Transform the response data - container_list: Final = ContainerListResponse(**response_data) + container_list: Final = ContainerListResponse.model_validate(raw_response.json()) return container_list @@ -235,7 +229,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request - data: Final[dict[str, object]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -245,9 +239,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: """Transform the OpenAI container retrieve response.""" - response_data: Final[OpenAIContainerPayload] = raw_response.json() - # Transform the response data - container_obj: Final = ContainerObject(**response_data) + container_obj: Final = ContainerObject.model_validate(raw_response.json()) return container_obj @@ -268,7 +260,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request - data: Final[dict[str, object]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -278,10 +270,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: """Transform the OpenAI container delete response.""" - response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json() - - # Transform the response data - delete_result: Final = DeleteContainerResult(**response_data) + delete_result: Final = DeleteContainerResult.model_validate(raw_response.json()) return delete_result @@ -326,10 +315,7 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: """Transform the OpenAI container file list response.""" - response_data: Final[OpenAIContainerFileListPayload] = raw_response.json() - - # Transform the response data - file_list: Final = ContainerFileListResponse(**response_data) + file_list: Final = ContainerFileListResponse.model_validate(raw_response.json()) return file_list @@ -352,7 +338,7 @@ class OpenAIContainerConfig(BaseContainerConfig): url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content") # No query parameters needed - params: Final[dict[str, object]] = {} + params: Final[dict[str, str]] = {} return url, params diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 16fa0017b23..1cfc6e06ee9 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -12,9 +12,14 @@ if TYPE_CHECKING: import openai from openai import AsyncOpenAI, OpenAI +from openai._base_client import make_request_options +from openai._constants import RAW_RESPONSE_HEADER +from openai._legacy_response import LegacyAPIResponse +from openai._types import RequestOptions +from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import overload import litellm @@ -43,6 +48,7 @@ from litellm.utils import ( from ...types.llms.openai import * from ..base import BaseLLM from .chat.gpt_5_transformation import OpenAIGPT5Config +from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, @@ -189,7 +195,12 @@ class OpenAIConfig(BaseConfig): elif litellm.openAIGPTAudioConfig.is_model_gpt_audio_model(model=model): return litellm.openAIGPTAudioConfig.get_supported_openai_params(model=model) else: - return litellm.openAIGPTConfig.get_supported_openai_params(model=model) + return self._gpt_config_for_model(model).get_supported_openai_params(model=model) + + def _gpt_config_for_model(self, model: str) -> OpenAIGPTConfig: + if type(self) is OpenAIConfig and not OpenAIGPTConfig.is_openai_catalog_model(model): + return OpenAIUnknownModelConfig() + return litellm.openAIGPTConfig def _map_openai_params(self, non_default_params: dict, optional_params: dict, model: str) -> dict: supported_openai_params: Final = self.get_supported_openai_params(model) @@ -231,7 +242,7 @@ class OpenAIConfig(BaseConfig): drop_params=drop_params, ) - return litellm.openAIGPTConfig.map_openai_params( + return self._gpt_config_for_model(model).map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, @@ -323,6 +334,28 @@ class OpenAIChatCompletionResponseIterator(BaseModelResponseIterator): raise e +_EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None) +_EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) +_NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) +_SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) + + +def _embedding_request_without_sdk_defaults( + data: Mapping[str, object], timeout: float | httpx.Timeout +) -> tuple[Mapping[str, object], RequestOptions]: + body: Final = { # mutable-ok: the SDK json-encodes the body and needs a plain dict + k: v for k, v in data.items() if k not in _SDK_OPTION_KEYS + } + extra_headers: Final = _EXTRA_HEADERS_ADAPTER.validate_python(data.get("extra_headers")) or _NO_EXTRA_HEADERS + options: Final = make_request_options( + extra_headers=types.MappingProxyType({**extra_headers, RAW_RESPONSE_HEADER: "true"}), + extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")), + extra_body=data.get("extra_body"), + timeout=timeout, + ) + return body, options + + class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): def __init__(self) -> None: super().__init__() @@ -1171,19 +1204,15 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> LegacyAPIResponse[CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = await openai_aclient.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return bypass_response + return await openai_aclient.embeddings.with_raw_response.create(**data, timeout=timeout) @track_llm_api_timing() def make_sync_openai_embedding_request( @@ -1192,20 +1221,15 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, - ): - """ - Helper to: - - call embeddings.create.with_raw_response when litellm.return_response_headers is True - - call embeddings.create by default - """ - try: - raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) - - headers: Final = dict(raw_response.headers) - response: Final = raw_response.parse() - return headers, response - except Exception as e: - raise e + ) -> LegacyAPIResponse[CreateEmbeddingResponse]: + if "encoding_format" not in data: + body, options = _embedding_request_without_sdk_defaults(data, timeout) + bypass_response: Final = openai_client.post( + "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse + ) + assert isinstance(bypass_response, LegacyAPIResponse) + return bypass_response + return openai_client.embeddings.with_raw_response.create(**data, timeout=timeout) async def aembedding( self, @@ -1230,14 +1254,15 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, shared_session=shared_session, ) - headers, response = await self.make_openai_embedding_request( + raw_response: Final = await self.make_openai_embedding_request( openai_aclient=openai_aclient, data=data, timeout=timeout, logging_obj=logging_obj, ) + headers: Final = dict(raw_response.headers) logging_obj.model_call_details["response_headers"] = headers - stringified_response: Final = response.model_dump() + stringified_response: Final = raw_response.parse().model_dump() ## LOGGING logging_obj.post_call( input=input, @@ -1329,13 +1354,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) ## embedding CALL - headers: dict | None = None - headers, sync_embedding_response = self.make_sync_openai_embedding_request( + raw_response: Final = self.make_sync_openai_embedding_request( openai_client=openai_client, data=data, timeout=timeout, logging_obj=logging_obj, ) + headers: Final = dict(raw_response.headers) + sync_embedding_response: Final = raw_response.parse() ## LOGGING logging_obj.model_call_details["response_headers"] = headers diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 24dbd7c08d0..1530c154e93 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,12 +28,16 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from collections.abc import Sequence +import time +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -41,17 +45,33 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_stream_usage, + stream_item_field, +) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ( AllMessageValues, + BaseLiteLLMOpenAIResponseObject, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ContentPartAddedEvent, + ContentPartDoneEvent, + ContentPartDonePartOutputText, ErrorEvent, ErrorEventError, OpenAIMcpServerTool, + OutputItemAddedEvent, + OutputItemDoneEvent, + OutputTextDeltaEvent, + OutputTextDoneEvent, + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, ResponsesAPIStreamEvents, + ResponsesAPIStreamingResponse, ) from litellm.types.responses.main import ( GenericResponseOutputItem, @@ -63,11 +83,13 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from fastapi import HTTPException - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseInputParam - from litellm.types.utils import ResponsesAPIResponse class ResponseOutputEnvelope(TypedDict, total=False): @@ -82,6 +104,10 @@ class ResponsesStreamChunk(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] + delta: ReadOnly[str] + item_id: ReadOnly[str] + output_index: ReadOnly[int] + content_index: ReadOnly[int] def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: @@ -658,8 +684,32 @@ class OpenAIResponsesHandler(BaseTranslation): def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. + + ``response.output_text.done`` events carry the whole part in ``text``, while + ``response.output_text.delta`` events carry fragments in ``delta``. A stream + that dies before its done event (``response.failed`` / ``response.incomplete``) + has text only in deltas, so per content part the done text wins when present + and the joined deltas fill in otherwise, never both. """ - return "".join([response.get("text", "") for response in responses_so_far]) + keyed_events: Final = tuple( + ( + (event.get("item_id"), event.get("output_index"), event.get("content_index")), + event.get("text"), + event.get("delta"), + ) + for event in responses_so_far + if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + ) + + def part_text(part_key: tuple[object, object, object]) -> str: + done_texts: Final = tuple( + text for key, text, _ in keyed_events if key == part_key and isinstance(text, str) + ) + if done_texts: + return done_texts[-1] + return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str)) + + return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events)) def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ @@ -837,3 +887,331 @@ class OpenAIResponsesHandler(BaseTranslation): content[content_idx]["text"] = guardrail_response elif hasattr(content[content_idx], "text"): content[content_idx].text = guardrail_response + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + """ + Build Responses API SSE events that deliver the guardrail block message + and terminate the stream cleanly, mirroring the non-streaming block + response: a completed response whose only output is the violation text, + with the real usage the upstream call consumed. + + - ``stream_started`` False (buffered / pre-stream): nothing has been + sent, so emit the full synthetic sequence (``response.created`` + through ``response.completed``). + - ``stream_started`` True (sampling / mid-stream): events already + reached the client, so continue the in-progress response: close the + output item still open on the wire, deliver the block message as a + new output item under the same response id, and close with a + ``response.completed`` carrying only the replacement item. + + The proxy's data generator appends ``data: [DONE]`` itself. + """ + events: Final = ( + self._block_continuation_events(exc, responses_so_far or ()) + if stream_started + else self._standalone_block_events(exc) + ) + return tuple( + f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode() + for event in events + ) + + @staticmethod + def _standalone_block_events(exc: "ModifyResponseException") -> Sequence[ResponsesAPIStreamingResponse]: + from litellm.responses.streaming_iterator import build_synthetic_response_events + + return build_synthetic_response_events( + transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model), + logging_obj=None, + chunk_size=max(len(exc.message), 1), + ) + + @staticmethod + def _block_continuation_events( + exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> Sequence[ResponsesAPIStreamingResponse]: + response_id, model, output_index = _continuation_identity(exc, responses_so_far) + item: Final = _blocked_output_item(exc) + item_id: Final = item.id + part: Final[_BlockedContentPart] = {"type": "output_text", "text": exc.message, "annotations": ()} + done_part: Final[_BlockedDoneContentPart] = { + "type": "output_text", + "text": exc.message, + "annotations": (), + "logprobs": None, + } + return ( + *_open_item_closing_events(responses_so_far), + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=item, + ), + ContentPartAddedEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id=item_id, + output_index=output_index, + content_index=0, + part=BaseLiteLLMOpenAIResponseObject.model_validate(part), + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + content_index=0, + delta=exc.message, + ), + OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=item_id, + output_index=output_index, + content_index=0, + text=exc.message, + ), + ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=item_id, + output_index=output_index, + content_index=0, + part=ContentPartDonePartOutputText.model_validate(done_part), + ), + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + item=item, + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=_blocked_response(exc, response_id=response_id, model=model, output_item=item), + ), + ) + + +class _BlockedContentPart(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + annotations: ReadOnly[tuple[object, ...]] + + +class _BlockedDoneContentPart(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + annotations: ReadOnly[tuple[object, ...]] + logprobs: ReadOnly[None] + + +class _BlockedItemPayload(TypedDict): + type: ReadOnly[str] + id: ReadOnly[str] + status: ReadOnly[str] + role: ReadOnly[str] + content: ReadOnly[tuple[_BlockedContentPart, ...]] + + +class _BlockedResponsePayload(TypedDict): + id: ReadOnly[str] + object: ReadOnly[str] + created_at: ReadOnly[int] + model: ReadOnly[str] + output: ReadOnly[tuple[GenericResponseOutputItem, ...]] + status: ReadOnly[str] + usage: ReadOnly[ResponseAPIUsage] + + +def _blocked_output_item(exc: "ModifyResponseException") -> GenericResponseOutputItem: + payload: Final[_BlockedItemPayload] = { + "type": "message", + "id": f"msg_{uuid.uuid4()}", + "status": "completed", + "role": "assistant", + "content": ({"type": "output_text", "text": exc.message, "annotations": ()},), + } + return GenericResponseOutputItem.model_validate(payload) + + +def _blocked_response( + exc: "ModifyResponseException", + response_id: str, + model: str, + output_item: GenericResponseOutputItem | None = None, +) -> ResponsesAPIResponse: + payload: Final[_BlockedResponsePayload] = { + "id": response_id, + "object": "response", + "created_at": int(time.time()), + "model": model, + "output": (output_item if output_item is not None else _blocked_output_item(exc),), + "status": "completed", + "usage": blocked_responses_stream_usage(exc.original_response), + } + return ResponsesAPIResponse.model_validate(payload) + + +def _continuation_identity(exc: "ModifyResponseException", responses_so_far: Sequence[object]) -> tuple[str, str, int]: + responses: Final = tuple( + response for item in responses_so_far if (response := stream_item_field(item, "response")) is not None + ) + response_id: Final = next( + (rid for response in responses if isinstance(rid := stream_item_field(response, "id"), str) and rid), + f"resp_{uuid.uuid4()}", + ) + model: Final = next( + (m for response in responses if isinstance(m := stream_item_field(response, "model"), str) and m), + exc.model, + ) + indices: Final = tuple( + index for item in responses_so_far if isinstance(index := stream_item_field(item, "output_index"), int) + ) + return response_id, model, max(indices) + 1 if indices else 0 + + +@dataclass(frozen=True, slots=True) +class _OpenItemState: + item_id: str + item_type: str + role: str + output_index: int + content_index: int + text: str + part_open: bool + payload: object + + +def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None: + typed: Final = tuple((stream_item_field(event, "type"), event) for event in responses_so_far) + added: Final = tuple( + (added_index, stream_item_field(event, "item")) + for event_type, event in typed + if event_type == "response.output_item.added" + and isinstance(added_index := stream_item_field(event, "output_index"), int) + ) + done_indices: Final = frozenset( + done_index + for event_type, event in typed + if event_type == "response.output_item.done" + and isinstance(done_index := stream_item_field(event, "output_index"), int) + ) + open_added: Final = tuple((index, payload) for index, payload in added if index not in done_indices) + if not open_added: + return None + output_index, item_payload = open_added[-1] + if item_payload is None: + return None + item_id: Final = stream_item_field(item_payload, "id") + if not isinstance(item_id, str) or not item_id: + return None + raw_type: Final = stream_item_field(item_payload, "type") + raw_role: Final = stream_item_field(item_payload, "role") + part_added: Final = tuple( + part_index + for event_type, event in typed + if event_type == "response.content_part.added" + and stream_item_field(event, "item_id") == item_id + and isinstance(part_index := stream_item_field(event, "content_index"), int) + ) + part_done: Final = frozenset( + part_done_index + for event_type, event in typed + if event_type == "response.content_part.done" + and stream_item_field(event, "item_id") == item_id + and isinstance(part_done_index := stream_item_field(event, "content_index"), int) + ) + open_parts: Final = tuple(index for index in part_added if index not in part_done) + text: Final = "".join( + delta + for event_type, event in typed + if event_type == "response.output_text.delta" + and stream_item_field(event, "item_id") == item_id + and isinstance(delta := stream_item_field(event, "delta"), str) + ) + return _OpenItemState( + item_id=item_id, + item_type=raw_type if isinstance(raw_type, str) and raw_type else "message", + role=raw_role if isinstance(raw_role, str) and raw_role else "assistant", + output_index=output_index, + content_index=open_parts[-1] if open_parts else 0, + text=text, + part_open=bool(open_parts), + payload=item_payload, + ) + + +_item_fields_adapter: Final = TypeAdapter(Mapping[str, object]) +_no_item_fields: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _incomplete_item_fields(payload: object) -> Mapping[str, object]: + raw: Final = payload.model_dump() if isinstance(payload, BaseModel) else payload + if not isinstance(raw, dict): + return _no_item_fields + return _item_fields_adapter.validate_python(raw) + + +def _open_item_closing_events(responses_so_far: Sequence[object]) -> Sequence[ResponsesAPIStreamingResponse]: + """Close the output item still in progress on the relayed stream before the + block item is appended: strict Responses clients reject a + ``response.completed`` that arrives while an earlier ``output_item.added`` + was never closed. A message item closes ``completed`` with exactly the text + the client has received so far; any other item type (a function call the + guardrail rejected, for instance) closes ``incomplete`` so the synthetic + done event can never authorize acting on it.""" + open_item: Final = _open_item_state(responses_so_far) + if open_item is None: + return () + if open_item.item_type != "message": + return ( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=open_item.output_index, + item=BaseLiteLLMOpenAIResponseObject.model_validate( + MappingProxyType({**_incomplete_item_fields(open_item.payload), "status": "incomplete"}) + ), + ), + ) + partial_part: Final[_BlockedContentPart] = { + "type": "output_text", + "text": open_item.text, + "annotations": (), + } + closed_payload: Final[_BlockedItemPayload] = { + "type": open_item.item_type, + "id": open_item.item_id, + "status": "completed", + "role": open_item.role, + "content": (partial_part,), + } + item_done: Final = OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=open_item.output_index, + item=GenericResponseOutputItem.model_validate(closed_payload), + ) + if not open_item.part_open: + return (item_done,) + partial_done_part: Final[_BlockedDoneContentPart] = { + "type": "output_text", + "text": open_item.text, + "annotations": (), + "logprobs": None, + } + return ( + OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=open_item.item_id, + output_index=open_item.output_index, + content_index=open_item.content_index, + text=open_item.text, + ), + ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=open_item.item_id, + output_index=open_item.output_index, + content_index=open_item.content_index, + part=ContentPartDonePartOutputText.model_validate(partial_done_part), + ), + item_done, + ) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index eadc087383a..01313e95878 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,10 +1,11 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints import httpx from openai.types.responses import ResponseReasoningItem from pydantic import BaseModel, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -14,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import * from litellm.types.responses.main import * @@ -35,6 +37,37 @@ else: _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) +_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) + + +class _DeleteResponseBody(TypedDict): + """Decoded body of the Responses API delete call.""" + + id: ReadOnly[str | None] + object: ReadOnly[str | None] + deleted: ReadOnly[bool | None] + + +class _DeleteResponse(Protocol): + """The delete call's HTTP response, read for the decoded body it carries.""" + + def json(self) -> _DeleteResponseBody: ... + + +class _JsonObjectResponse(Protocol): + """A Responses API HTTP response, read for the JSON object it decodes to.""" + + def json(self) -> dict[str, object]: ... + + +def _delete_response_body(response: _DeleteResponse) -> _DeleteResponseBody: + """Decode a delete response body into the id, object and deleted fields it carries.""" + return response.json() + + +def _json_object_body(response: _JsonObjectResponse) -> dict[str, object]: + """Decode a Responses API response body into its JSON object form.""" + return response.json() class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): @@ -179,8 +212,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools + replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input) final_request_params: Final = dict( - ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params) + ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) ) return final_request_params @@ -217,6 +251,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return input, tools + def _drop_foreign_tool_call_item_ids(self, input: str | ResponseInputParam) -> str | ResponseInputParam: + if self.custom_llm_provider not in _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS or not isinstance(input, list): + return input + sanitized_items: Final = [self._without_foreign_tool_call_item_id(item) for item in input] + return cast("ResponseInputParam", sanitized_items) # cast-ok: items keep their shape, minus a rejected id + + @staticmethod + def _without_foreign_tool_call_item_id(item: object) -> object: + if not isinstance(item, dict): + return item + item_type: Final = item.get("type") + item_id: Final = item.get("id") + genuine_prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type) if isinstance(item_type, str) else None + if genuine_prefix is None or not isinstance(item_id, str) or item_id.startswith(genuine_prefix): + return item + return {key: value for key, value in item.items() if key != "id"} # mutable-ok: outgoing JSON request item + def _flatten_tool_schema_combinators_for_openai( self, model: str, @@ -469,7 +520,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return None @staticmethod - def get_event_model_class(event_type: str) -> Any: + def get_event_model_class(event_type: str) -> type[BaseLiteLLMOpenAIResponseObject]: """ Returns the appropriate event model class based on the event type. @@ -583,7 +634,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the delete response API response into a DeleteResponseResult """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _delete_response_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) return DeleteResponseResult(**raw_response_json) @@ -618,7 +669,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the get response API response into a ResponsesAPIResponse """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers: Final = dict(raw_response.headers) @@ -646,7 +697,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}/input_items" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if before is not None: @@ -665,7 +716,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> dict: try: - return raw_response.json() + return _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) @@ -699,7 +750,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): Transform the cancel response API response into a ResponsesAPIResponse """ try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final = _json_object_body(raw_response) except Exception: raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code) raw_response_headers: Final = dict(raw_response.headers) @@ -742,7 +793,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) if sanitized_tools is not None: response_api_optional_request_params["tools"] = sanitized_tools - data: Final = dict(ResponsesAPIRequestParams(model=model, input=input, **response_api_optional_request_params)) + replay_safe_input: Final = self._drop_foreign_tool_call_item_ids(input) + data: Final = dict( + ResponsesAPIRequestParams(model=model, input=replay_safe_input, **response_api_optional_request_params) + ) return url, data diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 8c548b6b0d6..855c49c320b 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -5,10 +5,11 @@ For handling OpenAI-like chat completions, like IBM WatsonX, etc. """ import json -from collections.abc import Callable -from typing import Any, Final +from collections.abc import Callable, Mapping, Sequence +from typing import Final, TypedDict import httpx +from typing_extensions import ReadOnly import litellm from litellm import LlmProviders @@ -25,6 +26,23 @@ from ..common_utils import OpenAILikeBase, OpenAILikeError from .transformation import OpenAILikeChatConfig +class _OpenAILikeChatCompletion(TypedDict, total=False): + """The chat-completion JSON body an OpenAI-like provider returns for a non-streamed call.""" + + id: ReadOnly[str] + choices: ReadOnly[Sequence[Mapping[str, object]]] + created: ReadOnly[int] + model: ReadOnly[str] + system_fingerprint: ReadOnly[str] + usage: ReadOnly[Mapping[str, object]] + object: ReadOnly[str] + + +def _fake_streamed_model_response(payload: _OpenAILikeChatCompletion) -> ModelResponse: + """Build the single response a fake-streamed provider call replays as one chunk.""" + return ModelResponse(**payload) + + async def make_call( client: AsyncHTTPHandler | None, api_base: str, @@ -42,9 +60,9 @@ async def make_call( response: Final = await client.post(api_base, headers=headers, data=data, stream=not fake_stream) if streaming_decoder is not None: - completion_stream: Any = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) + completion_stream = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = ModelResponse(**response.json()) + model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) @@ -82,7 +100,7 @@ def make_sync_call( if streaming_decoder is not None: completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) elif fake_stream: - model_response: Final = ModelResponse(**response.json()) + model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) diff --git a/litellm/llms/parallel_ai/search/cost_calculator.py b/litellm/llms/parallel_ai/search/cost_calculator.py new file mode 100644 index 00000000000..809cd280cc8 --- /dev/null +++ b/litellm/llms/parallel_ai/search/cost_calculator.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.utils import get_model_info + +PARALLEL_AI_DEFAULT_RESULTS: Final = 10 +PARALLEL_AI_ADDITIONAL_RESULT_COST: Final = 0.001 +PARALLEL_AI_USAGE_PARAM: Final = "_parallel_ai_usage" +PARALLEL_AI_STANDARD_SEARCH_MODEL: Final = "parallel_ai/search" +PARALLEL_AI_FAST_SEARCH_MODEL: Final = "parallel_ai/search-fast" +PARALLEL_AI_TURBO_SEARCH_MODEL: Final = "parallel_ai/search-turbo" +PARALLEL_AI_PRICING_MODEL_BY_MODE: Final[Mapping[str, str]] = MappingProxyType( + { + "fast": PARALLEL_AI_FAST_SEARCH_MODEL, + "turbo": PARALLEL_AI_TURBO_SEARCH_MODEL, + } +) +ADVANCED_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _non_negative_int(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + return None + return value + + +def _usage_count(usage: Sequence[Mapping[str, object]], sku: str) -> int | None: + counts: Final = tuple( + count + for item in usage + if item.get("name") == sku + if (count := _non_negative_int(item.get("count"))) is not None + ) + return sum(counts) if counts else None + + +def _effective_mode(optional_params: Mapping[str, object]) -> str: + mode: Final = optional_params.get("mode") + if isinstance(mode, str): + return mode + + processor: Final = optional_params.get("processor") + if processor == "pro": + return "advanced" + return "basic" + + +def _effective_max_results(optional_params: Mapping[str, object]) -> int: + try: + advanced_settings: Final = ADVANCED_SETTINGS_ADAPTER.validate_python(optional_params.get("advanced_settings")) + advanced_max_results: Final = _non_negative_int(advanced_settings.get("max_results")) + if advanced_max_results is not None: + return advanced_max_results + except ValidationError: + pass + + max_results: Final = _non_negative_int(optional_params.get("max_results")) + return max_results if max_results is not None else PARALLEL_AI_DEFAULT_RESULTS + + +def _request_cost(mode: str) -> float: + pricing_model: Final = PARALLEL_AI_PRICING_MODEL_BY_MODE.get(mode, PARALLEL_AI_STANDARD_SEARCH_MODEL) + model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="parallel_ai") + return float(model_info.get("input_cost_per_query") or 0.0) + + +def _additional_results( + optional_params: Mapping[str, object], + usage: Sequence[Mapping[str, object]] | None, +) -> int: + usage_count: Final = _usage_count(usage, "sku_search_additional_results") if usage is not None else None + if usage_count is not None: + return usage_count + if usage is not None: + return 0 + return max(_effective_max_results(optional_params) - PARALLEL_AI_DEFAULT_RESULTS, 0) + + +def parallel_ai_search_cost( + optional_params: Mapping[str, object], + usage: Sequence[Mapping[str, object]] | None, +) -> float: + request_cost: Final = _request_cost(_effective_mode(optional_params)) + request_count_from_usage: Final = _usage_count(usage, "sku_search") if usage is not None else None + request_count: Final = request_count_from_usage if request_count_from_usage is not None else 1 + additional_results: Final = _additional_results(optional_params, usage) + return request_count * request_cost + additional_results * PARALLEL_AI_ADDITIONAL_RESULT_COST diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index ea21d1153fe..bde7b7b86db 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -4,9 +4,13 @@ Calls Parallel AI's /v1/search endpoint to search the web. Parallel AI API Reference: https://docs.parallel.ai/api-reference/search/search """ +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final, TypedDict import httpx +from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -14,9 +18,29 @@ from litellm.llms.base_llm.search.transformation import ( SearchResponse, SearchResult, ) +from litellm.llms.parallel_ai.search.cost_calculator import PARALLEL_AI_USAGE_PARAM from litellm.secret_managers.main import get_secret_str +class _ParallelAIV1SearchResult(BaseModel): + model_config = ConfigDict(extra="ignore") + + url: str | None = None + title: str | None = None + publish_date: str | None = None + excerpts: Sequence[str] | None = None + + +class _ParallelAIV1SearchResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + search_id: str | None = None + session_id: str | None = None + results: Sequence[_ParallelAIV1SearchResult] = () + usage: Sequence[Mapping[str, object]] | None = None + warnings: Sequence[Mapping[str, object]] | None = None + + class _ParallelAISourcePolicy(TypedDict, total=False): include_domains: list[str] exclude_domains: list[str] @@ -27,10 +51,16 @@ class _ParallelAIExcerptSettings(TypedDict, total=False): max_chars_per_result: int +class _ParallelAIFetchPolicy(TypedDict, total=False): + max_age_seconds: ReadOnly[int] + timeout_seconds: ReadOnly[float] + disable_cache_fallback: ReadOnly[bool] + + class _ParallelAIAdvancedSettings(TypedDict, total=False): source_policy: _ParallelAISourcePolicy excerpt_settings: _ParallelAIExcerptSettings - fetch_policy: dict + fetch_policy: _ParallelAIFetchPolicy location: str max_results: int @@ -43,14 +73,14 @@ class ParallelAISearchRequest(TypedDict, total=False): search_queries: list[str] # Required - at least one keyword search query objective: str # Optional - natural-language description of search goal - mode: str # Optional - 'turbo', 'basic', or 'advanced' (default 'advanced') + mode: str # Optional - 'turbo', 'fast', 'basic', or 'advanced' (default 'advanced') max_chars_total: int # Optional - upper bound on total excerpt characters session_id: str # Optional - tracks calls across search/extract requests client_model: str # Optional - model consuming the results advanced_settings: _ParallelAIAdvancedSettings -LEGACY_PROCESSOR_TO_MODE: Final = {"base": "basic", "pro": "advanced"} +LEGACY_PROCESSOR_TO_MODE: Final = MappingProxyType({"base": "basic", "pro": "advanced"}) class ParallelAISearchConfig(BaseSearchConfig): @@ -67,16 +97,16 @@ class ParallelAISearchConfig(BaseSearchConfig): api_base: str | None = None, **kwargs, ) -> dict: - api_key = self.resolve_server_api_key( + resolved_api_key: Final = self.resolve_server_api_key( caller_api_key=api_key, caller_api_base=api_base, key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), base_env_var="PARALLEL_AI_API_BASE", default_api_base=self.PARALLEL_AI_API_BASE, ) - if not api_key: + if not resolved_api_key: raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") - headers["x-api-key"] = api_key + headers["x-api-key"] = resolved_api_key headers["Content-Type"] = "application/json" return headers @@ -87,13 +117,12 @@ class ParallelAISearchConfig(BaseSearchConfig): data: dict | list[dict] | None = None, **kwargs, ) -> str: - api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE + resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/search"): - api_base = f"{api_base.removesuffix('/v1')}/v1/search" - - return api_base + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1/search"): + return trimmed + return f"{trimmed.removesuffix('/v1')}/v1/search" def transform_search_request( self, @@ -109,14 +138,17 @@ class ParallelAISearchConfig(BaseSearchConfig): - If string: maps to `search_queries` (single item) and `objective` - If list: maps to `search_queries` (keyword queries) optional_params: Optional parameters for the request - - mode: Search mode ('turbo', 'basic', 'advanced'); defaults to 'basic' + - mode: Search mode ('turbo', 'fast', 'basic', 'advanced'); defaults to 'basic' - processor: Legacy v1beta param; 'base' maps to mode 'basic', 'pro' to 'advanced' - max_results: Maximum number of search results -> `advanced_settings.max_results` - - search_domain_filter: Domains to include -> `advanced_settings.source_policy.include_domains` + - search_domain_filter / include_domains: Domains to include -> `advanced_settings.source_policy.include_domains` - exclude_domains: Domains to exclude -> `advanced_settings.source_policy.exclude_domains` - - country: ISO 3166-1 alpha-2 code -> `advanced_settings.location` + - after_date: RFC 3339 date (YYYY-MM-DD) -> `advanced_settings.source_policy.after_date` + - country / location: ISO 3166-1 alpha-2 code -> `advanced_settings.location` - max_chars_per_result: -> `advanced_settings.excerpt_settings.max_chars_per_result` - - Any other params are passed through to the request body as-is + - fetch_policy: Cache vs live-fetch policy -> `advanced_settings.fetch_policy` + - Any other params (objective, max_chars_total, session_id, client_model, ...) + are passed through to the request body as-is Returns: Dict with request data following the v1 search request spec @@ -137,7 +169,7 @@ class ParallelAISearchConfig(BaseSearchConfig): mode = LEGACY_PROCESSOR_TO_MODE.get(processor, processor) # the v1 API defaults to 'advanced' when mode is omitted; default to 'basic' # instead to keep v1beta's default tier (processor 'base') and litellm's - # $0.004/query cost map entry for `parallel_ai/search` accurate + # cost map entry for `parallel_ai/search` accurate request_data["mode"] = mode or "basic" advanced_settings: Final[_ParallelAIAdvancedSettings] = {} @@ -148,17 +180,29 @@ class ParallelAISearchConfig(BaseSearchConfig): if "country" in params: advanced_settings["location"] = params.pop("country") + if "location" in params: + advanced_settings["location"] = params.pop("location") + if "max_chars_per_result" in params: advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")} + if "fetch_policy" in params: + advanced_settings["fetch_policy"] = params.pop("fetch_policy") + source_policy: Final[_ParallelAISourcePolicy] = {} if "search_domain_filter" in params: source_policy["include_domains"] = params.pop("search_domain_filter") + if "include_domains" in params: + source_policy["include_domains"] = params.pop("include_domains") + if "exclude_domains" in params: source_policy["exclude_domains"] = params.pop("exclude_domains") + if "after_date" in params: + source_policy["after_date"] = params.pop("after_date") + if source_policy: advanced_settings["source_policy"] = source_policy @@ -170,9 +214,11 @@ class ParallelAISearchConfig(BaseSearchConfig): # unified-spec param with no v1 equivalent params.pop("max_tokens_per_page", None) - result_data: Final[dict] = dict(request_data) - result_data.update(params) - return result_data + # reserved for the provider's own reported usage, which prices the request; + # a caller-supplied value would otherwise set its own cost + params.pop(PARALLEL_AI_USAGE_PARAM, None) + + return {**request_data, **params} def transform_search_response( self, @@ -186,26 +232,49 @@ class ParallelAISearchConfig(BaseSearchConfig): Parallel AI -> LiteLLM mappings: - results[].title -> SearchResult.title - results[].url -> SearchResult.url - - results[].excerpts (array) -> SearchResult.snippet (joined string) + - results[].excerpts (array) -> SearchResult.snippet (joined string); the raw + array is preserved as an extra `excerpts` field on each result - results[].publish_date -> SearchResult.date + - search_id / session_id / warnings are preserved as extra fields on the + response; usage is preserved as `parallel_usage` (the `usage` name is + reserved for LiteLLM's token-usage object) """ - response_json: Final = raw_response.json() + parsed: Final = _ParallelAIV1SearchResponse.model_validate(raw_response.json()) - results: Final = [] - for result in response_json.get("results", []): - excerpts = result.get("excerpts") or [] - snippet = " ... ".join(excerpts) if excerpts else "" + # written unconditionally: leaving a caller-supplied value in place when the + # provider reports no usage would let the caller price its own request + logging_obj.optional_params = { + **logging_obj.optional_params, + PARALLEL_AI_USAGE_PARAM: parsed.usage, + } - search_result = SearchResult( - title=result.get("title") or "", - url=result.get("url") or "", - snippet=snippet, - date=result.get("publish_date"), - last_updated=None, + results: Final = tuple( + SearchResult.model_validate( + MappingProxyType( + { + "title": result.title or "", + "url": result.url or "", + "snippet": " ... ".join(result.excerpts or ()), + "date": result.publish_date, + "last_updated": None, + "excerpts": result.excerpts or (), + } + ) ) - results.append(search_result) - - return SearchResponse( - results=results, - object="search", + for result in parsed.results ) + + extra_fields: Final = MappingProxyType( + { + key: value + for key, value in ( + ("search_id", parsed.search_id), + ("session_id", parsed.session_id), + ("parallel_usage", parsed.usage), + ("warnings", parsed.warnings), + ) + if value is not None + } + ) + + return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields})) diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index cde65addb65..5913709c8a0 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -1,8 +1,10 @@ import asyncio import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import ( @@ -29,6 +31,16 @@ else: LiteLLMLoggingObj = Any +class _RunwayMLTask(TypedDict, total=False): + """The RunwayML task payload returned by POST /v1/text_to_image and GET /v1/tasks/{id}.""" + + id: ReadOnly[str] + status: ReadOnly[str] + output: ReadOnly[Sequence[str | Mapping[str, str]]] + failure: ReadOnly[str] + failureCode: ReadOnly[str] + + class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): """ Configuration for RunwayML image generation models. @@ -80,7 +92,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): @staticmethod def _transform_runwayml_response_to_openai( - response_data: dict[str, Any], + response_data: _RunwayMLTask, model_response: ImageResponse, ) -> ImageResponse: """ @@ -155,7 +167,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): raise TimeoutError(f"RunwayML task polling timed out after {timeout_secs} seconds") @staticmethod - def _check_task_status(response_data: dict[str, Any]) -> str: + def _check_task_status(response_data: _RunwayMLTask) -> str: """ Check RunwayML task status from response. @@ -227,7 +239,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayMLTask = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -276,7 +288,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayMLTask = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -322,7 +334,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): } """ try: - response_data = raw_response.json() + response_data: _RunwayMLTask = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", @@ -382,7 +394,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): We need to poll the task until it completes (status SUCCEEDED) using async polling. """ try: - response_data = raw_response.json() + response_data: _RunwayMLTask = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error transforming image generation response: {e}", diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index 1da8f0c66f0..19e6d8ff494 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -6,10 +6,11 @@ Maps OpenAI TTS spec to RunwayML Text-to-Speech API import asyncio import time -from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final, Union +from collections.abc import Coroutine, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union import httpx +from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger @@ -31,6 +32,14 @@ else: HttpxBinaryResponseContent = Any +class _RunwayTtsTaskResponse(TypedDict, total=False): + id: ReadOnly[str] + status: ReadOnly[str] + output: ReadOnly[Sequence[object]] + failure: ReadOnly[str] + failureCode: ReadOnly[str] + + class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for RunwayML Text-to-Speech @@ -64,7 +73,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): litellm_params_dict: dict, logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout, - extra_headers: dict[str, Any] | None, + extra_headers: dict[str, object] | None, base_llm_http_handler: Any, aspeech: bool, api_base: str | None, @@ -72,7 +81,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): **kwargs: Any, ) -> Union[ "HttpxBinaryResponseContent", - Coroutine[Any, Any, "HttpxBinaryResponseContent"], + Coroutine[object, object, "HttpxBinaryResponseContent"], ]: """ Dispatch method to handle RunwayML TTS requests @@ -242,7 +251,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): raise TimeoutError(f"RunwayML TTS task polling timed out after {timeout_secs} seconds") @staticmethod - def _check_task_status(response_data: dict[str, Any]) -> str: + def _check_task_status(response_data: _RunwayTtsTaskResponse) -> str: """ Check RunwayML task status from response. @@ -314,7 +323,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayTtsTaskResponse = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -362,7 +371,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): response = await client.get(url=task_url, headers=headers) response.raise_for_status() - response_data = response.json() + response_data: _RunwayTtsTaskResponse = response.json() # Check task status status = self._check_task_status(response_data=response_data) @@ -453,7 +462,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -483,7 +492,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() verbose_logger.debug("RunwayML TTS polling complete, downloading audio") @@ -522,7 +531,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): from litellm.types.llms.openai import HttpxBinaryResponseContent try: - response_data: Final = raw_response.json() + response_data: Final[_RunwayTtsTaskResponse] = raw_response.json() except Exception as e: raise self.get_error_class( error_message=f"Error parsing RunwayML TTS response: {e}", @@ -552,7 +561,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) # Get the completed task data - task_data: Final = polled_response.json() + task_data: Final[_RunwayTtsTaskResponse] = polled_response.json() verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio") diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index d7743d4d337..a2a93b6114a 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -8,9 +8,10 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path from threading import Lock -from typing import Any, Final +from typing import Any, Final, Protocol import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -33,8 +34,8 @@ def _get_home() -> str: return os.getenv(HOME_PATH_ENV_VAR, DEFAULT_HOME_PATH) -def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any: - cur: Any = d +def _get_nested(d: object, path: Sequence[str]) -> object: + cur: object = d if isinstance(cur, str): # This shouldn't happen if service keys are pre-parsed correctly try: @@ -54,7 +55,7 @@ def _get_nested(d: dict[str, Any] | str, path: Sequence[str]) -> Any: return cur -def _load_json_env(var_name: str) -> dict[str, Any] | None: +def _load_json_env(var_name: str) -> dict[str, object] | None: raw: Final = os.environ.get(var_name) if not raw: return None @@ -64,7 +65,7 @@ def _load_json_env(var_name: str) -> dict[str, Any] | None: return None -def _str_or_none(value) -> str | None: +def _str_or_none(value: object) -> str | None: try: return str(value) if value is not None else None except Exception: @@ -124,7 +125,7 @@ CREDENTIAL_VALUES: Final[list[CredentialsValue]] = [ ] -def init_conf(profile: str | None = None) -> dict[str, Any]: +def init_conf(profile: str | None = None) -> dict[str, object]: """ Loads config JSON from: 1) $AICORE_CONFIG if set, otherwise @@ -191,7 +192,7 @@ def resolve_resource_group(sources: list[Source]) -> str | None: def _parse_service_key_once( service_key: str | dict | None, -) -> dict[str, Any] | None: +) -> dict[str, object] | None: """ Pre-parse service_key if it's a string to avoid repeated JSON parsing. @@ -348,8 +349,33 @@ def validate_credentials( ) +class _TokenBody(TypedDict): + """Decoded body of the SAP AI Core OAuth2 token response.""" + + access_token: ReadOnly[str] + expires_in: ReadOnly[NotRequired[int]] + + +class _TokenResponse(Protocol): + """The token endpoint's HTTP response, read for the decoded token body it carries.""" + + def json(self) -> _TokenBody: ... + + +def _bearer_token_and_expiry(response: _TokenResponse) -> tuple[str, datetime]: + """Read a token response into the Authorization header value and the token's absolute expiry.""" + payload: Final = response.json() + expires_in: Final = int(payload.get("expires_in", 3600)) + access_token: Final = payload["access_token"] + return f"Bearer {access_token}", datetime.now(timezone.utc) + timedelta(seconds=expires_in) + + def _request_token( - client_id: str, auth_url: str, timeout: float, cert_pair=None, client_secret=None + client_id: str, + auth_url: str, + timeout: float, + cert_pair: tuple[str, str] | None = None, + client_secret: str | None = None, ) -> tuple[str, datetime]: data: Final = {"grant_type": "client_credentials", "client_id": client_id} if client_secret: @@ -361,15 +387,10 @@ def _request_token( with httpx.Client(cert=cert_pair) as raw_client: handler = HTTPHandler(client=raw_client) resp = handler.post(auth_url, data=data, timeout=timeout) - payload = resp.json() - else: - handler = _get_httpx_client() - resp = handler.post(auth_url, data=data, timeout=timeout) - payload = resp.json() - access_token: Final = payload["access_token"] - expires_in: Final = int(payload.get("expires_in", 3600)) - expiry_date: Final = datetime.now(timezone.utc) + timedelta(seconds=expires_in) - return f"Bearer {access_token}", expiry_date + return _bearer_token_and_expiry(resp) + handler = _get_httpx_client() + resp = handler.post(auth_url, data=data, timeout=timeout) + return _bearer_token_and_expiry(resp) except Exception as e: msg: Final = resp.text if resp is not None else getattr(e, "text", str(e)) raise RuntimeError(f"Token request failed: {msg}") from e diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 1de2337d8eb..a36c920dda0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,7 +1,7 @@ import re from copy import deepcopy from enum import Enum -from typing import Any, Final, Literal, get_type_hints +from typing import Any, Final, Literal, cast, get_type_hints import httpx @@ -31,7 +31,7 @@ class VertexAIError(BaseLLMException): super().__init__(message=message, status_code=status_code, headers=headers) -def redact_vertex_ai_metadata_from_logged_object(obj: Any) -> None: +def redact_vertex_ai_metadata_from_logged_object(obj: object) -> None: if isinstance(obj, dict): for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: if field in obj: @@ -651,7 +651,7 @@ def _build_json_schema(parameters: dict) -> dict: return parameters -def _filter_anyof_fields(schema_dict: dict[str, Any]) -> dict[str, Any]: +def _filter_anyof_fields(schema_dict: dict[str, object]) -> dict[str, object]: """ When anyof is present, only keep the anyof field and its contents - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 Filter out other fields in the same dict. @@ -704,7 +704,7 @@ def process_items(schema, depth=0): process_items(item, depth + 1) -def set_schema_property_ordering(schema: dict[str, Any], depth: int = 0) -> dict[str, Any]: +def set_schema_property_ordering(schema: dict[str, object], depth: int = 0) -> dict[str, object]: """ vertex ai and generativeai apis order output of fields alphabetically, unless you specify the order. python dicts retain order, so we just use that. Note that this field only applies to structured outputs, and not tools. @@ -724,14 +724,16 @@ def set_schema_property_ordering(schema: dict[str, Any], depth: int = 0) -> dict # retain propertyOrdering as an escape hatch if user already specifies it if "propertyOrdering" not in schema: schema["propertyOrdering"] = [k for k, v in schema["properties"].items()] - for k, v in schema["properties"].items(): - set_schema_property_ordering(v, depth + 1) - if "items" in schema: - set_schema_property_ordering(schema["items"], depth + 1) + for v in schema["properties"].values(): + if isinstance(v, dict): + set_schema_property_ordering(cast("dict[str, object]", v), depth + 1) # cast-ok: JSON Schema child + items: Final = schema.get("items") + if isinstance(items, dict): + set_schema_property_ordering(cast("dict[str, object]", items), depth + 1) # cast-ok: JSON Schema child return schema -def filter_schema_fields(schema_dict: dict[str, Any], valid_fields: set[str], processed=None) -> dict[str, Any]: +def filter_schema_fields(schema_dict: dict[str, object], valid_fields: set[str], processed=None) -> dict[str, object]: """ Recursively filter a schema dictionary to keep only valid fields. """ @@ -905,7 +907,7 @@ def _convert_schema_types(schema, depth=0): "maxProperties", } - any_of: Final[list[dict[str, Any]]] = [] + any_of: Final[list[dict[str, object]]] = [] for t in type_val: if not isinstance(t, str): continue @@ -916,7 +918,7 @@ def _convert_schema_types(schema, depth=0): # For object/array types, include type-specific fields if t in ("object", "array"): - item_schema = {"type": t} + item_schema: dict[str, object] = {"type": t} # Move type-specific fields into this anyOf item for field in type_specific_fields: if field in schema: @@ -1110,11 +1112,11 @@ class VertexAITokenCounter(BaseTokenCounter): self, model_to_use: str, messages: list[dict[str, Any]] | None, - contents: list[dict[str, Any]] | None, + contents: list[dict[str, object]] | None, deployment: dict[str, Any] | None = None, request_model: str = "", - tools: list[dict[str, Any]] | None = None, - system: Any | None = None, + tools: list[dict[str, object]] | None = None, + system: object | None = None, ) -> TokenCountResponse | None: import copy @@ -1131,25 +1133,26 @@ class VertexAITokenCounter(BaseTokenCounter): partner_models_handler: Final = VertexAIPartnerModels() # Extract vertex-specific params from litellm_params - vertex_project = count_tokens_params_request.get("vertex_project") or count_tokens_params_request.get( + partner_litellm_params: Final[dict[str, object]] = count_tokens_params_request + vertex_project = partner_litellm_params.get("vertex_project") or partner_litellm_params.get( "vertex_ai_project" ) - vertex_location = count_tokens_params_request.get("vertex_location") or count_tokens_params_request.get( + vertex_location = partner_litellm_params.get("vertex_location") or partner_litellm_params.get( "vertex_ai_location" ) # Count tokens not available on global location: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens - vertex_location = count_tokens_params_request.get("vertex_count_tokens_location") or vertex_location + vertex_location = partner_litellm_params.get("vertex_count_tokens_location") or vertex_location - vertex_credentials: Final = count_tokens_params_request.get( - "vertex_credentials" - ) or count_tokens_params_request.get("vertex_ai_credentials") + vertex_credentials: Final = partner_litellm_params.get("vertex_credentials") or partner_litellm_params.get( + "vertex_ai_credentials" + ) result = await partner_models_handler.count_tokens( model=model_to_use, messages=messages or [], - litellm_params=count_tokens_params_request, + litellm_params=partner_litellm_params, vertex_project=vertex_project, vertex_location=vertex_location, vertex_credentials=vertex_credentials, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b7f91bfba0d..b6ad9fbcc04 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -12,7 +12,7 @@ from urllib.parse import quote, unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, Required import litellm from litellm._uuid import uuid @@ -104,6 +104,27 @@ class _VertexBatchRow(TypedDict, total=False): processed_time: ReadOnly[str] +class _VertexEmbeddingVector(TypedDict): + values: ReadOnly[list[float]] + + +class _VertexEmbeddingUsageMetadata(TypedDict, total=False): + promptTokenCount: ReadOnly[int] + + +class _VertexEmbeddingResponse(TypedDict, total=False): + embedding: ReadOnly[Required[_VertexEmbeddingVector]] + usageMetadata: ReadOnly[_VertexEmbeddingUsageMetadata] + tokenCount: ReadOnly[int] + + +class _VertexEmbeddingBatchRow(TypedDict, total=False): + key: ReadOnly[str] + request: ReadOnly[Mapping[str, object]] + status: ReadOnly[Required[str]] + response: ReadOnly[Required[_VertexEmbeddingResponse]] + + class _OpenAIBatchOutputError(TypedDict): code: ReadOnly[str] message: ReadOnly[str] @@ -111,7 +132,7 @@ class _OpenAIBatchOutputError(TypedDict): class _OpenAIBatchOutputResponse(TypedDict): status_code: ReadOnly[int] - request_id: ReadOnly[str] + request_id: ReadOnly[object] body: ReadOnly[Mapping[str, object]] @@ -218,7 +239,7 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None return str(labels.get("litellm_custom_id", "unknown")) -def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: """ Whether a Vertex batch output row came from an `EmbedContentRequest`. @@ -237,7 +258,7 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) def _openai_batch_output_row( custom_id: str, - body: Mapping[str, Any] | None = None, + body: Mapping[str, object] | None = None, error_code: str | None = None, error_message: str = "", ) -> _OpenAIBatchOutputRow: @@ -259,7 +280,7 @@ def _openai_batch_output_row( } -def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int, int]: +def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]: """ Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch output row. @@ -278,7 +299,7 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) -def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: +def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> int: """ Prompt tokens billed for one Vertex Gemini Embedding batch row. @@ -293,7 +314,7 @@ def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, - vertex_output_rows: tuple[Mapping[str, Any], ...], + vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...], element_indices: tuple[int, ...], element_count: int, model: str | None, @@ -348,7 +369,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( def _transform_vertex_embeddings_batch_output_to_openai( - vertex_output_rows: Iterable[Mapping[str, Any]], + vertex_output_rows: Iterable[_VertexEmbeddingBatchRow], model: str | None, ) -> tuple[_OpenAIBatchOutputRow, ...]: """ @@ -388,7 +409,7 @@ def _model_from_managed_gcs_url(url: str) -> str | None: return match.group(1) if match else None -def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: +def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: """ Whether an OpenAI batch JSONL line targets the embeddings endpoint. @@ -431,7 +452,7 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" -def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, object]) -> Mapping[str, object]: """ One Vertex Gemini Embedding batch input row. @@ -453,8 +474,8 @@ def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( - openai_entry: Mapping[str, Any], -) -> tuple[Mapping[str, Any], ...]: + openai_entry: Mapping[str, object], +) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding batch rows, one per requested embedding. @@ -512,7 +533,7 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> tuple[Mapping[str, Any], ...]: +) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. @@ -533,7 +554,7 @@ def _openai_batch_jsonl_entry_to_vertex_rows( cached_content=None, ) - custom_id: Final = openai_entry.get("custom_id") + custom_id: Final[object] = openai_entry.get("custom_id") if custom_id is not None: if "labels" not in vertex_request_body: vertex_request_body["labels"] = {} diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 11c026010ee..e2d62be6a69 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -250,7 +250,7 @@ def _gs_uri_requires_content_type_metadata(url: str) -> bool: def _image_url_payload_may_need_sync_gcs_metadata_fetch( - raw_image_url: Any, + raw_image_url: object, ) -> bool: """ True when this image_url value (content-part image_url or assistant ``images[]`` @@ -326,7 +326,7 @@ def _openai_messages_may_need_sync_gcs_metadata_fetch( def _get_gcs_object_content_type( image_url: str, vertex_project: str | None = None, - vertex_credentials: Any | None = None, + vertex_credentials: object = None, ) -> str | None: """ Resolve content type from GCS object metadata. @@ -479,7 +479,7 @@ def _process_gemini_media( model: str | None = None, video_metadata: dict[str, Any] | None = None, vertex_project: str | None = None, - vertex_credentials: Any | None = None, + vertex_credentials: object = None, ) -> PartType: """ Given a media URL (image, audio, or video), return the appropriate PartType for Gemini @@ -1002,7 +1002,7 @@ def _gemini_convert_messages_with_history( if isinstance(_ss_invocations, list): for invocation in _ss_invocations: # Re-inject toolCall part - tc_part: dict[str, Any] = { + tc_part: dict[str, object] = { "toolCall": { "toolType": invocation.get("tool_type"), "id": invocation.get("id"), @@ -1015,13 +1015,13 @@ def _gemini_convert_messages_with_history( # Re-inject toolResponse part if response is present if "response" in invocation: - tr_dict: dict[str, Any] = { + tr_dict: dict[str, object] = { "id": invocation.get("id"), "response": invocation.get("response"), } if invocation.get("tool_type"): tr_dict["toolType"] = invocation["tool_type"] - tr_part: dict[str, Any] = {"toolResponse": tr_dict} + tr_part: dict[str, object] = {"toolResponse": tr_dict} if "response_thought_signature" in invocation: tr_part["thoughtSignature"] = invocation["response_thought_signature"] assistant_content.append(tr_part) @@ -1090,7 +1090,7 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: data_dict[k] = v -def _has_google_maps_tool(tools: Any | None) -> bool: +def _has_google_maps_tool(tools: object) -> bool: """Return True if any tool object in the list has a 'googleMaps' key.""" if not isinstance(tools, list): return False @@ -1127,7 +1127,7 @@ def _rewrite_mime_type_to_response_format(generation_config: GenerationConfig) - schema = generation_config.pop("response_schema", None) generation_config.pop("response_mime_type", None) - response_format: Final[dict[str, Any]] = {"text": {"mimeType": "APPLICATION_JSON"}} + response_format: Final[dict[str, dict[str, object]]] = {"text": {"mimeType": "APPLICATION_JSON"}} if schema is not None: response_format["text"]["schema"] = schema generation_config["responseFormat"] = response_format @@ -1316,7 +1316,7 @@ async def async_transform_request_body( timeout: float | httpx.Timeout | None, extra_headers: dict | None, optional_params: dict, - logging_obj: litellm.litellm_core_utils.litellm_logging.Logging, + logging_obj: LiteLLMLoggingObj, custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], litellm_params: dict, vertex_project: str | None, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d8b1e7ba17c..69fe5678de9 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -949,7 +949,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # For Gemini 3+ models, use thinkingLevel instead of thinkingBudget if model and VertexGeminiConfig._is_gemini_3_or_newer(model): if thinking_enabled: - if thinking_budget is None or thinking_budget == 0: + if thinking_budget == 0: params["includeThoughts"] = False else: params["includeThoughts"] = True diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index 2603552152d..b57a87c3325 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -177,8 +177,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): content_item = {"type": "image_url", "image_url": document_url} # Build DeepSeek OCR request + provider_model: Final = model if model.startswith("deepseek-ai/") else f"deepseek-ai/{model}" data: Final = { - "model": "deepseek-ai/" + model, + "model": provider_model, "messages": [{"role": "user", "content": [content_item]}], } diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index cf14ab88751..332f892ae6b 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -7,10 +7,14 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s import base64 from collections.abc import Coroutine +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union import httpx +from litellm.litellm_core_utils.audio_utils.utils import ( + speech_media_type_from_audio_bytes, +) from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -457,12 +461,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): if not response_content: raise ValueError("No audioContent in Vertex AI TTS response") - # Decode base64 to get binary content binary_data: Final = base64.b64decode(response_content) - - # Create an httpx.Response object with the binary data + media_type: Final = speech_media_type_from_audio_bytes(binary_data) response: Final = httpx.Response( status_code=200, + headers=None if media_type is None else MappingProxyType({"content-type": media_type}), content=binary_data, ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index d7ad69593c6..ef03e61a858 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -153,14 +153,17 @@ class VertexAIAnthropicConfig(AnthropicConfig): drop_params: bool, ) -> dict: """ - Override parent method to ensure VertexAI always uses tool-based structured outputs. - VertexAI doesn't support the output_format parameter, so we force all models - to use the tool-based approach for structured outputs. + Override parent method so VertexAI uses tool-based structured outputs + unless the vertex map entry advertises native structured output + (``output_format``, which Vertex AI Claude forwards for those models). """ - # Temporarily override model name to force tool-based approach - # This ensures Claude Sonnet 4.5 uses tools instead of output_format + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + original_model: Final = model - if "response_format" in non_default_params: + native_structured_output: Final = AnthropicModelInfo._get_provider_resolved_capability( + model, "supports_native_structured_output", "vertex_ai" + ) + if "response_format" in non_default_params and native_structured_output is not True: model = "claude-3-sonnet-20240229" # Use a model that will use tool-based approach # Call parent method with potentially modified model name diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index aca257dc095..1942bc850f1 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -9,7 +9,7 @@ import json import os import threading from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import urlparse import litellm @@ -47,6 +47,21 @@ else: GoogleCredentialsObject = Any +class _VertexCredentialsObject(Protocol): + """Structural view of the google-auth credentials handle that this class caches and refreshes.""" + + @property + def token(self) -> object: ... + + @property + def quota_project_id(self) -> str | None: ... + + @property + def expired(self) -> object: ... + + def refresh(self, request: object) -> None: ... + + class VertexBase: def __init__(self) -> None: super().__init__() @@ -55,7 +70,7 @@ class VertexBase: self._credentials: GoogleCredentialsObject | None = None self._credentials_project_mapping: dict[ tuple[VERTEX_CREDENTIALS_TYPES | None, str | None], - tuple[GoogleCredentialsObject, str | None], + tuple[_VertexCredentialsObject, str | None], ] = {} self.project_id: str | None = None self.async_handler: AsyncHTTPHandler | None = None @@ -109,7 +124,7 @@ class VertexBase: self, credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, - ) -> tuple[Any, str]: + ) -> tuple[_VertexCredentialsObject | None, str]: if credentials is not None: if isinstance(credentials, str): _is_path: Final = os.path.exists( @@ -209,7 +224,7 @@ class VertexBase: return creds, project_id # Google Auth Helpers -- extracted for mocking purposes in tests - def _credentials_from_identity_pool(self, json_obj, scopes): + def _credentials_from_identity_pool(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import identity_pool except ImportError: @@ -220,7 +235,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_pluggable(self, json_obj, scopes): + def _credentials_from_pluggable(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import pluggable except ImportError: @@ -231,7 +246,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_identity_pool_with_aws(self, json_obj, scopes): + def _credentials_from_identity_pool_with_aws(self, json_obj, scopes) -> _VertexCredentialsObject: try: from google.auth import aws except ImportError: @@ -242,7 +257,7 @@ class VertexBase: creds = creds.with_scopes(scopes) return creds - def _credentials_from_authorized_user(self, json_obj, scopes): + def _credentials_from_authorized_user(self, json_obj, scopes) -> _VertexCredentialsObject: try: import google.oauth2.credentials except ImportError: @@ -250,7 +265,7 @@ class VertexBase: return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes) - def _credentials_from_service_account(self, json_obj, scopes): + def _credentials_from_service_account(self, json_obj, scopes) -> _VertexCredentialsObject: try: import google.oauth2.service_account except ImportError: @@ -258,7 +273,7 @@ class VertexBase: return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes) - def _credentials_from_default_auth(self, scopes): + def _credentials_from_default_auth(self, scopes) -> tuple[_VertexCredentialsObject, str | None]: try: import google.auth as google_auth except ImportError: @@ -350,7 +365,7 @@ class VertexBase: ) return api_base - def refresh_auth(self, credentials: Any) -> None: + def refresh_auth(self, credentials: _VertexCredentialsObject) -> None: try: from google.auth.transport.requests import ( Request, @@ -426,7 +441,7 @@ class VertexBase: self, credential_cache_key: tuple, project_id: str | None, - ) -> tuple[str, str, "TokenState", Any, str | None] | None: + ) -> tuple[str, str, "TokenState", _VertexCredentialsObject, str | None] | None: """ Look up cached credentials and return usable token info for FRESH or STALE tokens (both are still valid for outbound requests). STALE @@ -449,7 +464,9 @@ class VertexBase: return None return creds.token, resolved_project, token_state, creds, cached_project_id - def _unpack_cached_credentials(self, credential_cache_key: tuple) -> tuple[Any, str | None]: + def _unpack_cached_credentials( + self, credential_cache_key: tuple + ) -> tuple[_VertexCredentialsObject | None, str | None]: """ Return (credentials, project_id) from the cache, or (None, None) if not cached. Handles both tuple and legacy cache formats. @@ -461,7 +478,7 @@ class VertexBase: return cached_entry return cached_entry, cached_entry.quota_project_id or getattr(cached_entry, "project_id", None) - def _get_token_state(self, credentials: Any) -> "TokenState": + def _get_token_state(self, credentials: _VertexCredentialsObject) -> "TokenState": """ Return the token state using google-auth's TokenState enum. @@ -485,7 +502,7 @@ class VertexBase: credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, credential_cache_key: tuple, - ) -> tuple[Any, str | None]: + ) -> tuple[_VertexCredentialsObject, str | None]: """Load credentials via load_auth (in thread) and cache the result.""" try: _credentials, credential_project_id = await asyncify(self.load_auth)( @@ -505,7 +522,7 @@ class VertexBase: async def _background_refresh_credentials( self, - credentials: Any, + credentials: _VertexCredentialsObject, credential_cache_key: tuple, credential_project_id: str | None, ) -> None: @@ -557,7 +574,7 @@ class VertexBase: def _schedule_background_refresh( self, - credentials: Any, + credentials: _VertexCredentialsObject, credential_cache_key: tuple, credential_project_id: str | None, ) -> None: @@ -575,7 +592,7 @@ class VertexBase: self._background_refresh_credentials(credentials, credential_cache_key, credential_project_id) ) - def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None: + def _drop_background_refresh_task(_fut: asyncio.Future[None]) -> None: if self._background_refresh_tasks.get(credential_cache_key) is _fut: self._background_refresh_tasks.pop(credential_cache_key, None) @@ -888,7 +905,7 @@ class VertexBase: # Convert dict credentials to string for caching cache_credentials: Final = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key: Final = (cache_credentials, project_id) - _credentials: GoogleCredentialsObject | None = None + _credentials: _VertexCredentialsObject | None = None verbose_logger.debug("Checking cached credentials for project_id: %s", project_id) diff --git a/litellm/main.py b/litellm/main.py index 0c8bff16f81..01c106adc7c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6292,18 +6292,15 @@ def embedding( if headers is not None and headers != {}: optional_params["extra_headers"] = headers - if encoding_format is not None: - optional_params["encoding_format"] = encoding_format + requested_encoding_format: Final = ( + encoding_format + or optional_params.get("encoding_format") + or get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") + ) + if requested_encoding_format is None or requested_encoding_format.strip().lower() == "none": + optional_params.pop("encoding_format", None) else: - env_fmt: Final = get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") - if env_fmt is not None and env_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - _default_fmt: Final = optional_params.get("encoding_format") or env_fmt or "float" - if _default_fmt.strip().lower() == "none": - optional_params.pop("encoding_format", None) - else: - optional_params["encoding_format"] = _default_fmt + optional_params["encoding_format"] = requested_encoding_format api_version = None @@ -6952,12 +6949,18 @@ def embedding( aembedding=aembedding, headers=headers, ) - elif custom_llm_provider == "dashscope": - dashscope_key: Final = api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY") + elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"): + from litellm.llms.dashscope.common_utils import ( + missing_dashscope_family_key_message, + resolve_dashscope_family_api_key, + ) + + dashscope_key: Final = resolve_dashscope_family_api_key( + custom_llm_provider=custom_llm_provider, + api_key=api_key or litellm.api_key, + ) if dashscope_key is None: - raise ValueError( - "Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter." - ) + raise ValueError(missing_dashscope_family_key_message(custom_llm_provider)) if extra_headers is not None and isinstance(extra_headers, dict): headers = extra_headers else: @@ -8634,6 +8637,16 @@ def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Opti hidden_params["response_cost"] = response_cost +def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: + if logging_obj is None: + return + if isinstance(getattr(usage, "cost", None), (int, float)): + return + computed_cost: Final = logging_obj._response_cost_calculator(result=response) + if isinstance(computed_cost, (int, float)) and computed_cost > 0: + setattr(usage, "cost", computed_cost) + + def stream_chunk_builder( chunks: list, messages: list | None = None, @@ -8728,12 +8741,7 @@ def stream_chunk_builder( ) break - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - setattr( - usage, - "cost", - logging_obj._response_cost_calculator(result=response), - ) + _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) @@ -8912,10 +8920,7 @@ def stream_chunk_builder( ) break - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - setattr(usage, "cost", logging_obj._response_cost_calculator(result=response)) - + _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 718e6c489fd..10e8b56d925 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -554,6 +554,7 @@ "supports_vision": true }, "amazon.nova-sonic-v1:0": { + "deprecation_date": "2026-09-14", "input_cost_per_audio_token": 3.4e-06, "input_cost_per_token": 6e-08, "litellm_provider": "bedrock", @@ -1450,6 +1451,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -1487,6 +1526,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "global.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -1524,6 +1601,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "us.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 2.75e-07, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "eu.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -1561,6 +1676,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "eu.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 2.75e-07, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -1591,7 +1744,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1780,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1816,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1852,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1888,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1924,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2217,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2254,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2291,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2328,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2365,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2402,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -3045,6 +3198,7 @@ "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { + "deprecation_date": "2027-12-05", "supports_mid_conversation_system": true, "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -3077,7 +3231,43 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "azure_ai/claude-fable-5-1": { + "supports_mid_conversation_system": true, + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-5": { + "deprecation_date": "2027-07-08", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3110,6 +3300,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-4-8": { + "deprecation_date": "2027-09-01", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3188,6 +3379,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-5": { + "deprecation_date": "2027-06-30", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -9451,7 +9643,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Flash": { "input_cost_per_image_token": 1.75e-06, @@ -9464,7 +9657,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", @@ -9963,7 +10157,9 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.45e-07, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", @@ -9977,6 +10173,24 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.8e-08, + "supports_prompt_caching": true + }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 1.4e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 4.4e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { @@ -10192,11 +10406,13 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supports_function_calling": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1e-07, + "supports_prompt_caching": true }, "azure_ai/kimi-k2.6": { "deprecation_date": "2027-04-16", @@ -10207,7 +10423,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k2-6-in-microsoft-foundry/4513125", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supported_modalities": [ "text", "image" @@ -10218,7 +10434,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.6e-07, + "supports_prompt_caching": true }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, @@ -11902,7 +12120,7 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 2.65e-06, + "output_cost_per_token": 6e-07, "supports_pdf_input": true }, "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { @@ -12287,6 +12505,7 @@ "supports_tool_choice": true }, "cerebras/zai-glm-4.7": { + "deprecation_date": "2026-08-17", "input_cost_per_token": 2.25e-06, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -13022,6 +13241,47 @@ "supports_native_structured_output": true, "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, + "claude-fable-5-1": { + "deprecation_date": "2027-09-01", + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/models/fable-5-1/overview" + }, "claude-opus-5": { "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, @@ -14718,6 +14978,1910 @@ "/v1/images/generations" ] }, + "qwencloud/deepseek-v4-flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/deepseek-v4-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/glm-5.1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 202745, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/glm-5.2": { + "cache_read_input_token_cost": 2.8e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 229376, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen-coder": { + "input_cost_per_token": 3e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-flash": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-flash-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-max": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 30720, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 6.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-01-25": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-04-28": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-07-14": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-plus-2025-09-11": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-plus-latest": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-turbo": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-2024-11-01": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-2025-04-28": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-latest": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-30b-a3b": { + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-coder-flash": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-flash-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-plus-2025-07-22": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-max-preview": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-max": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-max-2026-01-23": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "qwencloud/qwen3.5-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3.7-max": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3.7-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3.8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwq-plus": { + "input_cost_per_token": 8e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 98304, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-image-2.0": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-2.0-pro": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-3.0": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-3.0-pro": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/deepseek-v4-flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/deepseek-v4-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/glm-5.1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 202745, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/glm-5.2": { + "cache_read_input_token_cost": 2.8e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 229376, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen-coder": { + "input_cost_per_token": 3e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-flash": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-flash-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-max": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 30720, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 6.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-01-25": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-04-28": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-07-14": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-plus-2025-09-11": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-plus-latest": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-turbo": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-2024-11-01": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-2025-04-28": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-latest": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-30b-a3b": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-coder-flash": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-flash-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-plus-2025-07-22": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max-preview": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max-2026-01-23": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.5-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.7-max": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3.7-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwq-plus": { + "input_cost_per_token": 8e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 98304, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-image-2.0": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-2.0-pro": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-3.0": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-3.0-pro": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "databricks/databricks-bge-large-en": { "cache_creation_input_token_cost": 1.0003e-07, "cache_read_input_token_cost": 1.0003e-07, @@ -15101,6 +17265,62 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/databricks-deepseek-v4-flash-0731": { + "cache_creation_input_token_cost": 1.4e-07, + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "input_dbu_cost_per_token": 2e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "output_dbu_cost_per_token": 4e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "databricks/databricks-deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 1.31999e-06, + "cache_read_input_token_cost": 1.3202e-07, + "input_cost_per_token": 1.31999e-06, + "input_dbu_cost_per_token": 1.8857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 3.95997e-06, + "output_dbu_cost_per_token": 5.6571e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "databricks/databricks-gemini-2-5-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, @@ -20631,6 +22851,7 @@ "supports_image_size": false }, "gemini-live-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -21303,6 +23524,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -23140,6 +25418,65 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -23548,6 +25885,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, @@ -23852,8 +26246,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23867,7 +26263,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23895,8 +26292,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23910,7 +26309,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -26881,16 +29281,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -26902,6 +29305,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -26944,16 +29348,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -26965,6 +29372,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -27008,16 +29416,19 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -27029,6 +29440,7 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -27071,16 +29483,19 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -27092,6 +29507,7 @@ "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, @@ -27331,7 +29747,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -27385,7 +29804,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -27534,7 +29956,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -27583,7 +30008,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -27632,7 +30060,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, @@ -27681,7 +30111,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -28617,17 +31049,18 @@ }, "gpt-realtime-2": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, "litellm_provider": "openai", - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, - "output_cost_per_token": 1.6e-05, + "output_cost_per_token": 2.4e-05, "supported_endpoints": [ "/v1/realtime" ], @@ -28691,8 +31124,8 @@ "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, @@ -28724,7 +31157,7 @@ "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "realtime", @@ -31488,19 +33921,21 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/magistral-medium-latest": { - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 5e-06, - "source": "https://mistral.ai/news/magistral", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-2506": { "deprecation_date": "2025-11-30", @@ -31519,19 +33954,21 @@ "supports_tool_choice": true }, "mistral/magistral-small-latest": { - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://mistral.ai/pricing#api-pricing", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-1-2-2509": { "deprecation_date": "2026-07-31", @@ -31663,16 +34100,21 @@ "supports_vision": true }, "mistral/mistral-medium": { - "input_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.1e-06, + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-medium-2312": { "deprecation_date": "2025-06-16", @@ -36339,12 +38781,22 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, "parallel_ai/search": { - "input_cost_per_query": 0.004, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-fast": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, "parallel_ai/search-pro": { - "input_cost_per_query": 0.009, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-turbo": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, @@ -39113,13 +41565,13 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6.25e-06, + "output_cost_per_token": 6e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -39729,6 +42181,70 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "us-gov.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us-gov.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_1hr": 2.2e-06, @@ -41892,6 +44408,42 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "vertex_ai/claude-fable-5-1": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", "regional_endpoint_uplift_multiplier": 1.1, @@ -41927,6 +44479,42 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "vertex_ai/claude-fable-5-1@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, @@ -43440,7 +46028,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43453,8 +46042,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43469,7 +46060,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43483,8 +46075,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43544,6 +46138,26 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "voyage/rerank-3": { + "input_cost_per_token": 5e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/rerank-3-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-2": { "input_cost_per_token": 1e-07, "litellm_provider": "voyage", @@ -44795,6 +47409,27 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-build-latest": { + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/developers/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, @@ -52231,14 +54866,14 @@ "supports_vision": true }, "fireworks_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 6.6e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -55115,6 +57750,34 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, @@ -55124,5 +57787,592 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "scaleway/glm-5.2": { + "input_cost_per_token": 1.8e-06, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false + }, + "scaleway/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": false + }, + "azure_ai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "deprecation_date": "2026-10-03", + "input_cost_per_token": 9.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-west-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { + "use_openai_responses_path": true, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://aws.amazon.com/bedrock/pricing/", + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2.4e-07 + }, + "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "azure/us-gov/gpt-5.1": { + "cache_read_input_token_cost": 1.71875e-07, + "default_reasoning_effort": "none", + "input_cost_per_token": 1.71875e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.375e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure/us-gov/o3-mini": { + "cache_read_input_token_cost": 7.57e-07, + "input_cost_per_token": 1.513e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.05e-06, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "azure/us-gov/text-embedding-3-large": { + "input_cost_per_token": 1.63e-07, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "azure/us-gov/text-embedding-3-small": { + "input_cost_per_token": 2.5e-08, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/openai/whisper": { + "input_cost_per_second": 7.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "cloudflare/@cf/openai/whisper-large-v3-turbo": { + "input_cost_per_second": 8.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/litellm/models/user.py b/litellm/models/user.py index 259c3440d87..82f78c28078 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -7,7 +7,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from from datetime import datetime -from pydantic import ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.organization_membership import ( @@ -67,3 +67,11 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): if not self.models: return True return model_name in self.models + + +class SCIMPlaceholder(BaseModel): + """A user row keyed by a value that names another account by SSO identity or email.""" + + placeholder_user_id: str + resolved_user_ids: tuple[str, ...] + team_ids: tuple[str, ...] diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9095cee15a9..c4bd03fb1c3 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -6,7 +6,7 @@ from __future__ import annotations import asyncio import contextvars -from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator, Iterator +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator from functools import partial from types import TracebackType from typing import Any, Final, cast @@ -27,19 +27,19 @@ base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils -async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, Any]: +async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]: async for chunk in iterable: yield chunk -def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, Any, Any]: +def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, bytes, None]: yield from iterable -class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): +class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): def __init__( self, - response: Coroutine[Any, Any, httpx.Response], + response: Awaitable[httpx.Response], litellm_logging_obj: LiteLLMLoggingObj, provider_config: BasePassthroughConfig, ) -> None: @@ -48,7 +48,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): self._headers = httpx.Headers() self._response_coro = response self._response: httpx.Response - self._iterator: AsyncGenerator[bytes, Any] + self._iterator: AsyncGenerator[bytes, bytes] self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks @@ -172,7 +172,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[Any, Any]): pass -class PassthroughStreamingResponse(Generator[Any, Any, Any]): +class PassthroughStreamingResponse(Generator[bytes, bytes, None]): def __init__( self, response: httpx.Response, @@ -184,7 +184,7 @@ class PassthroughStreamingResponse(Generator[Any, Any, Any]): self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj self._provider_config = provider_config - self._iterator: Generator[bytes, Any, Any] = _as_generator(response.iter_bytes()) + self._iterator: Generator[bytes, bytes, None] = _as_generator(response.iter_bytes()) self._raw_bytes: list[bytes] = [] # mutable-ok: instance buffer for streaming chunks self._flush_scheduled = False @@ -263,7 +263,7 @@ async def allm_passthrough_route( cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, **kwargs, -) -> httpx.Response | AsyncGenerator[Any, Any]: +) -> httpx.Response | AsyncGenerator[bytes, bytes]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -390,10 +390,10 @@ def llm_passthrough_route( **kwargs, ) -> ( httpx.Response - | Coroutine[Any, Any, httpx.Response] - | Coroutine[Any, Any, httpx.Response | AsyncGenerator[Any, Any]] - | Generator[Any, Any, Any] - | AsyncGenerator[Any, Any] + | Coroutine[object, object, httpx.Response] + | Coroutine[object, object, httpx.Response | AsyncGenerator[bytes, bytes]] + | Generator[bytes, bytes, None] + | AsyncGenerator[bytes, bytes] ): """ Pass through requests to the LLM APIs. @@ -592,7 +592,7 @@ async def _async_passthrough_request( is_streaming_request: bool, litellm_logging_obj: LiteLLMLoggingObj, provider_config: BasePassthroughConfig, -) -> httpx.Response | AsyncGenerator[Any, Any]: +) -> httpx.Response | AsyncGenerator[bytes, bytes]: """ Handle async passthrough requests. Uses async client to send request and properly handles streaming. diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index ead26ab65c5..9d6b1e18f59 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -671,6 +671,42 @@ "interactions": true } }, + "qwencloud": { + "display_name": "QwenCloud (`qwencloud`)", + "url": "https://docs.litellm.ai/docs/providers/qwencloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "qwen_ai_platform": { + "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "url": "https://docs.litellm.ai/docs/providers/qwencloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, "databricks": { "display_name": "Databricks (`databricks`)", "url": "https://docs.litellm.ai/docs/providers/databricks", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2330120adad..1f552ff3e13 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast @@ -1206,7 +1206,7 @@ def _deserialize_json_dict(data: str | _StringMap | None) -> dict[str, str] | No return data -def _deserialize_json_list(data: Any) -> list[dict[str, Any]] | None: +def _deserialize_json_list(data: object) -> list[dict[str, Any]] | None: """Deserialize a JSON array stored in the DB (``env_vars`` and friends). Returns ``None`` for empty / null / unparseable input. Accepts strings @@ -1219,7 +1219,7 @@ def _deserialize_json_list(data: Any) -> list[dict[str, Any]] | None: return None if isinstance(data, str): try: - parsed: Final = json.loads(data) + parsed: Final[object] = json.loads(data) except (json.JSONDecodeError, TypeError): return None data = parsed @@ -1914,7 +1914,7 @@ class MCPServerManager: async def load_servers_from_config( self, - mcp_servers_config: dict[str, Any], + mcp_servers_config: dict[str, MCPServerConfig], mcp_aliases: dict[str, str] | None = None, ): """ @@ -3068,7 +3068,7 @@ class MCPServerManager: return {} cache_key: Final = "toolset_perms:" + ",".join(sorted(toolset_ids)) - cached: Final = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final[dict[str, list[str]] | None] = await user_api_key_cache.async_get_cache(key=cache_key) if cached is not None: return cached @@ -5154,7 +5154,7 @@ class MCPServerManager: # Wrapped so the bridge runs inside the task: the caller only holds the task and # gathers it later, so there is no other point that still sees a block here. - async def _run_during_call_hook() -> Mapping[str, Any] | None: + async def _run_during_call_hook() -> Mapping[str, object] | None: try: return await proxy_logging_obj.during_call_hook( user_api_key_dict=user_api_key_auth, @@ -5656,7 +5656,7 @@ class MCPServerManager: async def _gather_openapi_tool_tasks( self, - tasks: list[Any], + tasks: Sequence[Awaitable[object]], proxy_logging_obj: ProxyLogging | None, ) -> CallToolResult: """Await OpenAPI tool tasks and return the tool call result.""" diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..d1ef73a15cd 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,13 +1,16 @@ import asyncio import importlib -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal +import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -18,8 +21,11 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListOk, + ServerOutcome, classify_list_exception, list_fault_http_status, + outcome_wire_value, ) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, @@ -86,8 +92,6 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: - from mcp.types import Tool as MCPTool - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -99,6 +103,7 @@ if MCP_AVAILABLE: ListMCPToolsRestAPIResponseObject, MCPInfo, MCPServer, + _aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes _apply_toolset_scope, _fire_mcp_tool_call_logging, execute_mcp_tool, @@ -803,9 +808,6 @@ if MCP_AVAILABLE: list(allowed_server_ids_set), _rest_client_ip ) - list_tools_result: Final = [] - error_message = None - # If server_id is specified, only query that specific server if server_id: return await _list_tools_for_single_server( @@ -849,22 +851,19 @@ if MCP_AVAILABLE: else {} ) - # Query all servers the user has access to - errors: Final = [] - for allowed_server_id in allowed_server_ids: - server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) - if server is None: - continue - - server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header) - user_oauth_extra_headers = await _get_user_oauth_extra_headers( + async def list_server( + server: MCPServer, + ) -> tuple[Sequence[ListMCPToolsRestAPIResponseObject], ServerOutcome]: + server_auth_header: Final = _get_server_auth_header( + server, mcp_server_auth_headers, mcp_auth_header + ) + user_oauth_extra_headers: Final = await _get_user_oauth_extra_headers( server, user_api_key_dict, prefetched_creds=prefetched_oauth_creds, ) - try: - tools_result = await _get_tools_for_single_server( + tools_result: Final = await _get_tools_for_single_server( server, server_auth_header, raw_headers_from_request, @@ -872,24 +871,36 @@ if MCP_AVAILABLE: extra_headers=user_oauth_extra_headers, apply_tool_filters=apply_tool_filters, ) - list_tools_result.extend(tools_result) except Exception as e: verbose_logger.exception("Error getting tools from %s: %s", server.name, e) - errors.append( - f"{get_server_prefix(server)}: {classify_list_exception(e).tag}" - if isinstance(e, (MCPServerListError, MCPUpstreamAuthError)) - else f"{get_server_prefix(server)}: {e}" - ) - continue + return (), classify_list_exception(e) + return tools_result, ServerListOk(tool_count=len(tools_result)) - if errors and not list_tools_result: - error_message = "Failed to get tools from servers: " + "; ".join(errors) - - return { - "tools": list_tools_result, - "error": "partial_failure" if error_message else None, - "message": (error_message if error_message else "Successfully retrieved tools"), - } + # Query all servers the user has access to + queried_servers: Final = tuple( + server + for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) + if server is not None + ) + listings: Final = tuple([await list_server(server) for server in queried_servers]) + list_tools_result: Final = [tool for tools, _ in listings for tool in tools] + server_outcomes: Final = MappingProxyType( + {_aggregate_server_key(server): outcome for server, (_, outcome) in zip(queried_servers, listings)} + ) + errors: Final = tuple( + f"{key}: {outcome.tag}" for key, outcome in server_outcomes.items() if outcome.tag != "ok" + ) + error_message: Final = ( + "Failed to get tools from servers: " + "; ".join(errors) + if errors and not list_tools_result + else None + ) + return { + "tools": list_tools_result, + "error": "partial_failure" if error_message else None, + "message": (error_message if error_message else "Successfully retrieved tools"), + "server_outcomes": {key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items()}, + } except MCPUpstreamAuthError as e: # Surface upstream pass-through 401/403 challenges to the client so @@ -1173,6 +1184,7 @@ if MCP_AVAILABLE: transport=request.transport, auth_type=request.auth_type, mcp_info=request.mcp_info, + timeout=request.timeout, command=request.command, args=request.args, env=request.env, @@ -1402,11 +1414,28 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - async def _list_tools_session_operation(session): - return await session.list_tools() - - list_tools_response: Final = await client.run_with_session(_list_tools_session_operation) - list_tools_result: Final[list[MCPTool]] = list_tools_response.tools + # Bound the whole pagination walk: without this the preview is limited only by the + # per-request timeout times the page cap. max() keeps the pre-pagination guarantee + # that a single slow page within the client timeout still succeeds, and a + # per-server timeout above the global default extends the deadline with it. + listing_deadline: Final = max( + getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, + MCP_TOOL_LISTING_TIMEOUT, + ) + list_tools_result = None # rebind-ok: set inside the timeout scope below + with anyio.move_on_after(listing_deadline): + list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + if list_tools_result is None: + verbose_logger.warning( + "MCP tools/list preview timed out after %s seconds while paginating upstream tools", + listing_deadline, + ) + return { # mutable-ok: error response payload + "status": "error", + "error": True, + "message": f"Timed out listing tools after {listing_deadline} seconds. " + "The MCP server may be responding slowly or paginating excessively.", + } model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 7ec0f4b5192..dcf1b01bc25 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -5,6 +5,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints. """ import asyncio +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger @@ -74,7 +75,7 @@ class SemanticMCPToolFilter: self.router_instance = litellm_router_instance self.tool_router: SemanticRouter | None = None self.context_window_error: str | None = None - self._tool_map: dict[str, Any] = {} # MCPTool objects or OpenAI function dicts + self._tool_map: dict[str, object] = {} # MCPTool objects or OpenAI function dicts self._index_sync_lock = asyncio.Lock() async def build_router_from_mcp_registry(self) -> None: @@ -182,11 +183,11 @@ class SemanticMCPToolFilter: return raise - def _has_tools_missing_from_index(self, tools: list[Any]) -> bool: + def _has_tools_missing_from_index(self, tools: Sequence[object]) -> bool: """Allocation-free check for any named tool not yet in the semantic index.""" return any(name and name not in self._tool_map for name in (self._extract_tool_info(t)[0] for t in tools)) - def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]: + def _tools_missing_from_index(self, tools: Sequence[object]) -> Mapping[str, object]: """Map name -> tool for every named tool not yet in the semantic index.""" return { name: tool @@ -194,7 +195,7 @@ class SemanticMCPToolFilter: if name and name not in self._tool_map } - async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None: + async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: """ Index request-time tools the startup build never saw. @@ -385,7 +386,7 @@ class SemanticMCPToolFilter: separator: Final = client_name[-len(canonical) - 1] return separator in ("_", "-") - def _get_tools_by_names(self, tool_names: list[str], available_tools: list[Any]) -> list[Any]: + def _get_tools_by_names(self, tool_names: Sequence[str], available_tools: Sequence[object]) -> list[object]: """ Get tools from available_tools by their names, preserving the semantic router's ordering. @@ -401,14 +402,14 @@ class SemanticMCPToolFilter: # Exact matches win over suffix matches when both are present, and # each incoming tool is returned at most once even if two canonical # names happen to be tail-compatible with the same incoming name. - available_by_name: Final[dict[str, Any]] = {} + available_by_name: Final[dict[str, object]] = {} for tool in available_tools: client_name, _ = self._extract_tool_info(tool) if client_name and client_name not in available_by_name: available_by_name[client_name] = tool - matched: Final[list[Any]] = [] - used_ids: Final[set] = set() + matched: Final[list[object]] = [] + used_ids: Final[set[int]] = set() for canonical in tool_names: tool = available_by_name.get(canonical) if tool is None: @@ -430,7 +431,7 @@ class SemanticMCPToolFilter: used_ids.add(id(tool)) return matched - def extract_user_query(self, messages: list[dict[str, Any]]) -> str: + def extract_user_query(self, messages: Sequence[Mapping[str, object]]) -> str: """ Extract user query from messages for /chat/completions or /responses. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5af45b29226..13c7a4c7cfa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -32002,6 +32002,62 @@ "title": "SCIMPatchOperation", "type": "object" }, + "SCIMPlaceholder": { + "description": "A user row keyed by a value that names another account by SSO identity or email.", + "properties": { + "placeholder_user_id": { + "title": "Placeholder User Id", + "type": "string" + }, + "resolved_user_ids": { + "items": { + "type": "string" + }, + "title": "Resolved User Ids", + "type": "array" + }, + "team_ids": { + "items": { + "type": "string" + }, + "title": "Team Ids", + "type": "array" + } + }, + "required": [ + "placeholder_user_id", + "resolved_user_ids", + "team_ids" + ], + "title": "SCIMPlaceholder", + "type": "object" + }, + "SCIMPlaceholderMergeResult": { + "properties": { + "merged_into_user_id": { + "title": "Merged Into User Id", + "type": "string" + }, + "placeholder_user_id": { + "title": "Placeholder User Id", + "type": "string" + }, + "team_ids": { + "items": { + "type": "string" + }, + "title": "Team Ids", + "type": "array" + } + }, + "required": [ + "placeholder_user_id", + "merged_into_user_id", + "team_ids" + ], + "title": "SCIMPlaceholderMergeResult", + "type": "object" + }, "SCIMServiceProviderConfig": { "properties": { "authenticationSchemes": { @@ -33641,6 +33697,129 @@ "scim" ] } + }, + "/scim/v2/placeholders": { + "get": { + "description": "List user rows whose id is another account's SSO identity or email.\n\nAn earlier release provisioned a group member it could not match as a user keyed\nby the raw member value, and that row now shadows the account the value really\nnames, so every push of that member is refused. This lists those rows so an\noperator can fold each one into the account it shadows with\n``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of\nits own or owns virtual keys is left out: someone uses that account.", + "operationId": "list_placeholders_scim_v2_placeholders_get", + "parameters": [ + { + "in": "query", + "name": "feature", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Feature" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/SCIMPlaceholder" + }, + "title": "Response List Placeholders Scim V2 Placeholders Get", + "type": "array" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Placeholders", + "tags": [ + "scim" + ] + } + }, + "/scim/v2/placeholders/{user_id}/merge": { + "post": { + "description": "Fold a placeholder user into the one account its id names by SSO identity or email.\n\nThe account is added to every team the placeholder is on, then the placeholder is\ndeleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group\npush resolves the member value to the real account. Refused with 409 when the row\nhas an SSO identity of its own, owns virtual keys, or names no account or several.", + "operationId": "merge_placeholder_scim_v2_placeholders__user_id__merge_post", + "parameters": [ + { + "in": "path", + "name": "user_id", + "required": true, + "schema": { + "title": "User ID", + "type": "string" + } + }, + { + "in": "query", + "name": "feature", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Feature" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SCIMPlaceholderMergeResult" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Merge Placeholder", + "tags": [ + "scim" + ] + } } } }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 428736da96d..5c0ba55d411 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,6 +1216,13 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None + @field_validator("team_id", mode="before") + @classmethod + def treat_cleared_team_id_as_unset(cls, v: object) -> object: + if v == "": + return None + return v + class GenerateKeyResponse(KeyRequestBase): key: str @@ -2436,9 +2443,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): database_socket_timeout: float | None = Field( None, description=( - "Prisma `socket_timeout` URL param (seconds). When set, an idle/slow " - "connection that has not produced data within this window is closed. " - "This is the main knob for capping idle DB connections from LiteLLM." + "Prisma `socket_timeout` URL param (seconds). When set, an in-flight " + "operation that has not produced data within this window is aborted. " + "For capping how long idle pooled connections are kept, see " + "`database_max_idle_connection_lifetime`." + ), + ) + database_max_idle_connection_lifetime: float | None = Field( + 60, + description=( + "Prisma `max_idle_connection_lifetime` URL param (seconds). A pooled " + "connection idle longer than this is closed and replaced instead of " + "being handed to the next request. Defaults to 60 so connections are " + "recycled before common infra idle timeouts (AWS NLB / RDS Proxy " + "~350s, many LBs 60-350s) silently drop them and requests fail with " + "`Error { kind: Closed }`. A value pinned on the DATABASE_URL or set " + "via `database_extra_connection_params` takes precedence." ), ) database_extra_connection_params: dict[str, Any] | None = Field( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index bd02cfdf907..31b05320cd3 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -14,7 +14,7 @@ import json from collections.abc import AsyncGenerator, Mapping from copy import deepcopy from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -215,11 +215,20 @@ def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None: ) +class _JsonRpcResponse(Protocol): + def json(self) -> dict[str, object]: ... + + +def _jsonrpc_body(response: _JsonRpcResponse) -> dict[str, object]: + """The decoded JSON-RPC body of ``response``.""" + return response.json() + + async def _forward_jsonrpc( agent_url: str, body: dict[str, object], extra_headers: Mapping[str, str] | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -230,7 +239,7 @@ async def _forward_jsonrpc( ) resp: Final = await handler.post(agent_url, json=body, headers=headers) try: - result: Final = resp.json() + result: Final = _jsonrpc_body(resp) except Exception: resp.raise_for_status() raise @@ -940,8 +949,8 @@ async def invoke_agent_a2a( ) result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers) if method == "agent/getAuthenticatedExtendedCard": - if isinstance(result.get("result"), dict): - card: Final = result["result"] + card: Final = result.get("result") + if isinstance(card, dict): proxy_url: Final = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}") # Rewrite the upstream agent URL in both 0.3 (top-level `url`) # and 1.0 (`supportedInterfaces[0].url`) wire formats so that diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py new file mode 100644 index 00000000000..7da5e5099fc --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -0,0 +1,174 @@ +""" +Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only +stream event carrying a model, so streamed responses report the requested model like +non-streaming ones do. + +Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the +provider passthrough path) or as event dicts (fake-stream and agentic paths). +""" + +import json +import re +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +_MESSAGE_START_EVENT: Final = "message_start" +_MESSAGE_START_MARKER: Final = b"message_start" +_SSE_DATA_FIELD: Final = "data:" +_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n") +_MAX_HELD_BYTES: Final = 65536 +_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') + +_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None: + message: Final = event.get("message") + if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict): + return None + if message.get("model") == requested_model: + return None + return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is + + +def _restamped_data_line(line: str, requested_model: str) -> str | None: + stripped: Final = line.strip() + if not stripped.startswith(_SSE_DATA_FIELD): + return None + payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip() + if not payload or payload == "[DONE]": + return None + try: + event: Final = _EVENT_ADAPTER.validate_json(payload) + except ValidationError: + return None + restamped: Final = _restamped_event(event, requested_model) + if restamped is None: + return None + terminator: Final = line[len(line.rstrip("\r\n")) :] + return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}" + + +def _restamped_frame(frame: str, requested_model: str) -> str | None: + lines: Final = frame.splitlines(keepends=True) + restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) + if all(line is None for line in restamped): + return None + return "".join(new if new is not None else old for new, old in zip(restamped, lines)) + + +def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: + """ + Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``. + + Chunks that carry no model are returned unchanged. + """ + if isinstance(chunk, dict): + try: + event: Final = _EVENT_ADAPTER.validate_python(chunk) + except ValidationError: + return chunk + return _restamped_event(event, requested_model) or chunk + + if isinstance(chunk, (bytes, bytearray)): + if _MESSAGE_START_EVENT.encode() not in chunk: + return chunk + restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model) + return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8") + + if isinstance(chunk, str): + if _MESSAGE_START_EVENT not in chunk: + return chunk + restamped_text: Final = _restamped_frame(chunk, requested_model) + return chunk if restamped_text is None else restamped_text + + return chunk + + +def _is_ping_frame(frame: bytes) -> bool: + return any(marker in frame for marker in _PING_MARKERS) + + +class AnthropicStreamModelRestamper: + """ + Per-stream restamper for the encoded passthrough path, where chunks are raw + transport reads: the ``message_start`` SSE frame can arrive split across + chunks or coalesced with later frames. Complete frames (``\\n\\n``, + ``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator + closes them and an incomplete tail is held until it completes, so the + restamp never misses a torn frame; ``flush`` returns whatever is still held + when the stream ends so no bytes are swallowed. Once ``message_start`` has + been handled, or the first real event proves the stream carries none, every + later chunk passes through untouched. + """ + + def __init__(self, requested_model: str) -> None: + self._requested_model: Final = requested_model + self._held = b"" + self._armed = True + + def process(self, chunk: object) -> object: + if not self._armed: + return chunk + if isinstance(chunk, (bytes, bytearray)): + return self._process_encoded(bytes(chunk)) + if isinstance(chunk, str): + return self._process_encoded(chunk.encode("utf-8")) + restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model) + if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"): + self._armed = False + return restamped + + def flush(self) -> bytes: + held: Final = self._held + self._held = b"" + self._armed = False + if not held: + return b"" + restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model) + return restamped if isinstance(restamped, bytes) else held + + def _process_encoded(self, data: bytes) -> bytes: + combined: Final = self._held + data + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined)) + if not boundaries: + if len(combined) > _MAX_HELD_BYTES: + self._held = b"" + self._armed = False + return combined + self._held = combined + return b"" + emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]]) + tail: Final = combined[boundaries[-1] :] + if not self._armed: + self._held = b"" + return emitted + tail + self._held = tail + return emitted + + def _restamped_closed_block(self, closed: bytes) -> bytes: + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed)) + frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries)) + decider: Final = next( + ( + index + for index, frame in enumerate(frames) + if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame)) + ), + None, + ) + if decider is None: + return closed + self._armed = False + if _MESSAGE_START_MARKER not in frames[decider]: + return closed + restamped_text: Final = _restamped_frame( + frames[decider].decode("utf-8", errors="ignore"), self._requested_model + ) + if restamped_text is None: + return closed + return b"".join( + restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames) + ) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 83ee10b8108..89b2c92cdfd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -994,7 +994,7 @@ def get_key_model_rpm_limit( # 2. Check model_max_budget if user_api_key_dict.model_max_budget: - model_rpm_limit: Final[dict[str, Any]] = {} + model_rpm_limit: Final[dict[str, int]] = {} for model, budget in user_api_key_dict.model_max_budget.items(): if isinstance(budget, dict) and budget.get("rpm_limit") is not None: model_rpm_limit[model] = budget["rpm_limit"] @@ -1037,7 +1037,7 @@ def get_key_model_tpm_limit( # 2. Check model_max_budget (iterate per-model like RPM does) if user_api_key_dict.model_max_budget: - model_tpm_limit: Final[dict[str, Any]] = {} + model_tpm_limit: Final[dict[str, int]] = {} for model, budget in user_api_key_dict.model_max_budget.items(): if isinstance(budget, dict) and budget.get("tpm_limit") is not None: model_tpm_limit[model] = budget["tpm_limit"] @@ -1100,7 +1100,7 @@ def _validated_output_token_estimates_per_model(raw: object) -> Mapping[str, int def _estimated_output_tokens_from_metadata( - metadata: Mapping[str, Any] | None, + metadata: Mapping[str, object] | None, model_name: str | None, ) -> int | None: """Resolve the per-model, then global, estimate out of one metadata blob. @@ -1666,7 +1666,7 @@ def _dedupe_model_candidates(candidates: list[str]) -> list[str]: return deduped -def _get_case_insensitive_mapping_value(mapping: Mapping[str, Any] | None, key: str) -> Any: +def _get_case_insensitive_mapping_value(mapping: Mapping[str, object] | None, key: str) -> object: if not mapping: return None if key in mapping: @@ -1770,8 +1770,8 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non def _extract_model_candidates_from_request( request_data: dict, route: str, - request_headers: Mapping[str, Any] | None = None, - request_query_params: Mapping[str, Any] | None = None, + request_headers: Mapping[str, object] | None = None, + request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, ) -> list[str]: candidates: Final[list[str]] = [] @@ -1863,8 +1863,8 @@ def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool def get_model_from_request( request_data: dict, route: str, - request_headers: Mapping[str, Any] | None = None, - request_query_params: Mapping[str, Any] | None = None, + request_headers: Mapping[str, object] | None = None, + request_query_params: Mapping[str, object] | None = None, llm_router: Router | None = None, request: Request | None = None, ) -> str | list[str] | None: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 39e6ca9a369..0795cee7409 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,8 +14,8 @@ import hashlib import os import re import time -from collections.abc import Awaitable, Callable -from typing import Any, Final, Literal, NoReturn, TypeVar, cast +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast import httpx import jwt @@ -24,6 +24,7 @@ from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from fastapi import HTTPException, status from jwt.api_jwk import PyJWK +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value @@ -93,6 +94,47 @@ UNREACHABLE_CACHE_KEY_PREFIX: Final = "litellm_jwks_unreachable_" _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) +class _JWTAuthSettings(Protocol): + """The JWT auth settings block this handler reads back through ``getattr``, when one is configured.""" + + @property + def issuers(self) -> Sequence[JWTIssuerConfig] | None: ... + + @property + def public_key_ttl(self) -> float: ... + + @property + def public_key_stale_ttl(self) -> float: ... + + +class _OIDCDiscoveryBody(TypedDict, total=False): + """Decoded OIDC discovery document, read for the JWKS endpoint it advertises.""" + + jwks_uri: ReadOnly[str] + + +class _OIDCDiscoveryResponse(Protocol): + """The discovery endpoint's HTTP response, read for the decoded document it carries.""" + + def json(self) -> _OIDCDiscoveryBody: ... + + +class _UserInfoResponse(Protocol): + """The OIDC UserInfo endpoint's HTTP response, read for the identity document it carries.""" + + def json(self) -> dict[str, object]: ... + + +def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody: + """Decode an OIDC discovery response body.""" + return response.json() + + +def _userinfo_document(response: _UserInfoResponse) -> dict[str, object]: + """Decode an OIDC UserInfo response body into its JSON object form.""" + return response.json() + + def jwks_unavailable_exception(error: JWKSUnreachableError) -> ProxyException: return ProxyException( message=( @@ -794,7 +836,7 @@ class JWTHandler: f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}" ) try: - discovery: Final = response.json() + discovery: Final = _discovery_document(response) except Exception as e: raise Exception(f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}") @@ -806,13 +848,13 @@ class JWTHandler: return jwks_uri def _get_public_key_cache_ttl(self) -> float: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return 600 return litellm_jwtauth.public_key_ttl def _get_public_key_stale_ttl(self) -> float: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return DEFAULT_JWKS_STALE_TTL return litellm_jwtauth.public_key_stale_ttl @@ -938,7 +980,7 @@ class JWTHandler: if response.status_code != 200: raise Exception(f"OIDC UserInfo endpoint returned status {response.status_code}: {response.text}") - userinfo: Final = response.json() + userinfo: Final = _userinfo_document(response) verbose_proxy_logger.debug("Received OIDC UserInfo: %s", userinfo) # Cache the userinfo response @@ -996,7 +1038,7 @@ class JWTHandler: } def _get_configured_issuer(self, token: str) -> JWTIssuerConfig | None: - litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + litellm_jwtauth: Final[_JWTAuthSettings | None] = getattr(self, "litellm_jwtauth", None) if litellm_jwtauth is None: return None diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 45e05d353fb..c591cbabee1 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -9,6 +9,7 @@ import click import requests from .auth import context_secret_vault, get_stored_api_key, login +from .cmd_quoting import quote_for_cmd ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" @@ -151,31 +152,9 @@ def verify_proxy_key( _WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"}) -_CMD_PERCENT_GUARD: Final = "%%cd:~,%" _CMD_LINE_BREAKS: Final = ("\r", "\n") -def _double_trailing_backslashes(segment: str) -> str: - bare: Final = segment.rstrip("\\") - return bare + "\\" * 2 * (len(segment) - len(bare)) - - -def _quote_for_cmd(token: str) -> str: - """Quote one token so both parsers that read it see the original text. - - Follows the algorithm the Rust standard library settled on for batch files - after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a - quoted string on a lone `"` and so wants an embedded one doubled, and the - shim's own interpreter, which re-splits `%*` under C runtime rules where a - backslash escapes the quote that follows it, so every backslash run standing - before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each - `%` is prefixed with `%%cd:~,`: the zero-length substring of the always - defined `cd` expands to nothing and leaves no `%` pair for cmd to match. - """ - escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"')) - return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"' - - def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: """Build what CreateProcess runs, routing batch shims through cmd.exe. @@ -202,7 +181,7 @@ def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on " "Windows: cmd.exe ends the command line there, so the agent would silently lose it." ) - inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest)) + inner: Final = " ".join(quote_for_cmd(token) for token in (path, *rest)) return f'cmd.exe /d /e:on /v:off /s /c "{inner}"' diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index ea1fa019c83..46af641636e 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -8,6 +8,7 @@ live here rather than in either command module. import shlex import shutil +import sys from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -17,6 +18,8 @@ from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.private_json import write_private_json +from .cmd_quoting import quote_for_cmd + ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" @@ -87,9 +90,12 @@ def merge_claude_settings( return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} -def resolve_api_key_helper(base_url: str) -> str: +def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: """Build the shell command Claude Code should run for its apiKeyHelper. + Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe + on Windows, so every token is quoted for the shell that will read it. + Resolves `lite` to an absolute path so the helper works regardless of the PATH visible to whatever subprocess Claude Code spawns it from. Passing --base-url explicitly (rather than relying on the bare invocation Claude @@ -106,7 +112,8 @@ def resolve_api_key_helper(base_url: str) -> str: raise ClaudeSettingsError( "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it." ) - return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token" + quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote + return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: diff --git a/litellm/proxy/client/cli/commands/cmd_quoting.py b/litellm/proxy/client/cli/commands/cmd_quoting.py new file mode 100644 index 00000000000..efd6d584527 --- /dev/null +++ b/litellm/proxy/client/cli/commands/cmd_quoting.py @@ -0,0 +1,26 @@ +"""Quoting for command lines that cmd.exe reads before handing them to a program.""" + +from typing import Final + +_CMD_PERCENT_GUARD: Final = "%%cd:~,%" + + +def _double_trailing_backslashes(segment: str) -> str: + bare: Final = segment.rstrip("\\") + return bare + "\\" * 2 * (len(segment) - len(bare)) + + +def quote_for_cmd(token: str) -> str: + """Quote one token so both parsers that read it see the original text. + + Follows the algorithm the Rust standard library settled on for batch files + after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a + quoted string on a lone `"` and so wants an embedded one doubled, and the + program's own C runtime argv split, where a backslash escapes the quote that + follows it, so every backslash run standing before a quote is doubled. + Quoting cannot stop cmd expanding `%VAR%`, so each `%` is prefixed with + `%%cd:~,`: the zero-length substring of the always defined `cd` expands to + nothing and leaves no `%` pair for cmd to match. + """ + escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"')) + return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"' diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1d860b02875..05ddef822f1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -176,6 +176,9 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, +) from litellm.proxy.litellm_pre_call_utils import ( add_litellm_data_to_request, refresh_proxy_server_request_body_snapshot, @@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, proxy_logging_obj=proxy_logging_obj, request=request, + restamp_model=( + None if _should_return_raw_model_name(self.data) else requested_model_from_client + ), ) return await create_response( generator=wrap_sse_stream_with_keepalive_pings( @@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing: else: return chunk + @staticmethod + def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer: + if restamper is None: + return ProxyBaseLLMRequestProcessing.return_sse_chunk + + def serialize(chunk: object) -> str: + return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) + + return serialize + @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, @@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing: serialize_chunk: StreamChunkSerializer, serialize_error: StreamErrorSerializer, request: Request | None = None, + flush_tail: Callable[[], bytes] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, cost injection, then yields chunks via serialize_chunk; on exception runs failure hook and yields via serialize_error. Use for SSE or NDJSON. + + ``flush_tail`` runs once after the upstream iterator completes cleanly and + its non-empty result is yielded, so a serializer that buffers bytes across + chunks can emit anything still held at end of stream. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing: # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) + held_tail: Final = flush_tail() if flush_tail is not None else b"" + if held_tail: + yield serialize_chunk(held_tail) stream_completed = True except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit @@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing: # billing and release exactly once. This is the outermost generator # Starlette closes on disconnect, so the nested iterator hook (which # only sees GeneratorExit on GC) cannot own the refund. - if not stream_completed: - client_disconnected = True + client_disconnected = not stream_completed if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, @@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing: request_data: dict, proxy_logging_obj: ProxyLogging, request: Request | None = None, + restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events. @@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing: SSE serializers directly (rather than re-wrapping it in another ``async for: yield`` trampoline), so a streamed chunk traverses one fewer async-generator layer / coroutine resume on the hot path. + + ``restamp_model`` publishes that name on the Anthropic ``message_start`` + event in place of the provider's model, matching what the non-streaming + response reports. """ + restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), request=request, + flush_tail=None if restamper is None else restamper.flush, ) @overload diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a28d03eebf0..39e74d2c8bd 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -533,8 +533,8 @@ def sanitize_openai_provider_metadata( Strips LiteLLM proxy-internal tracking fields that must not be forwarded to OpenAI batch/file APIs. """ - if not metadata: - return metadata + if metadata is None: + return None sanitized: Final[dict[str, str]] = {} for key, value in metadata.items(): if key in LITELLM_PROXY_INTERNAL_METADATA_KEYS: @@ -547,7 +547,7 @@ def sanitize_openai_provider_metadata( key, type(value).__name__, ) - return sanitized or None + return None if metadata and not sanitized else sanitized def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_name: str | None): @@ -650,7 +650,7 @@ def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object] return [c.lower() if isinstance(c, str) else c for c in callbacks] -def strip_callback_config(metadata: dict[str, Any] | None) -> dict[str, Any] | None: +def strip_callback_config(metadata: dict[str, object] | None) -> dict[str, object] | None: """Return key/team metadata without the slots that carry callback credentials.""" if not isinstance(metadata, dict): return metadata diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index bc7b80801fe..2a20e7b07ce 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -1,8 +1,12 @@ from collections.abc import Mapping, Sequence -from typing import Any, Final +from typing import Final, TypeAlias, Union from litellm._logging import verbose_proxy_logger +JsonValue: TypeAlias = Union["JsonObject", "JsonArray", str, int, float, bool, None] +JsonObject: TypeAlias = dict[str, JsonValue] +JsonArray: TypeAlias = list[JsonValue] + class CustomOpenAPISpec: """ @@ -27,7 +31,20 @@ class CustomOpenAPISpec: RESPONSES_API_PATHS = ["/v1/responses", "/responses"] @staticmethod - def get_pydantic_schema(model_class) -> Mapping[str, object] | None: + def _as_object(node: JsonValue) -> JsonObject: + return node if isinstance(node, dict) else {} + + @staticmethod + def _as_array(node: JsonValue) -> JsonArray: + return node if isinstance(node, list) else [] + + @staticmethod + def _components_schemas(openapi_schema: JsonObject) -> JsonObject: + components: Final = CustomOpenAPISpec._as_object(openapi_schema.setdefault("components", {})) + return CustomOpenAPISpec._as_object(components.setdefault("schemas", {})) + + @staticmethod + def get_pydantic_schema(model_class) -> JsonObject | None: """ Get JSON schema from a Pydantic model, handling both v1 and v2 APIs. @@ -54,9 +71,7 @@ class CustomOpenAPISpec: return None @staticmethod - def add_schema_to_components( - openapi_schema: dict[str, Any], schema_name: str, schema_def: Mapping[str, object] - ) -> None: + def add_schema_to_components(openapi_schema: JsonObject, schema_name: str, schema_def: JsonObject) -> None: """ Add a schema definition to the OpenAPI components/schemas section. @@ -66,16 +81,25 @@ class CustomOpenAPISpec: schema_def: The schema definition """ # Ensure components/schemas structure exists - if "components" not in openapi_schema: - openapi_schema["components"] = {} - if "schemas" not in openapi_schema["components"]: - openapi_schema["components"]["schemas"] = {} + _ = CustomOpenAPISpec._components_schemas(openapi_schema) # Add the schema CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) @staticmethod - def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: Sequence[str], schema_ref: str) -> None: + def _expanded_request_field(field_name: str, field_def: JsonValue) -> JsonValue: + expanded: Final = CustomOpenAPISpec._rewrite_defs_refs( + CustomOpenAPISpec._expand_field_definition(CustomOpenAPISpec._as_object(field_def)) + ) + if field_name != "messages": + return expanded + return { + **CustomOpenAPISpec._as_object(expanded), + "example": [{"role": "user", "content": "Hello, how are you?"}], + } + + @staticmethod + def add_request_body_to_paths(openapi_schema: JsonObject, paths: Sequence[str], schema_ref: str) -> None: """ Add request body with expanded form fields for better Swagger UI display. This keeps the request body but expands it to show individual fields in the UI. @@ -86,54 +110,58 @@ class CustomOpenAPISpec: schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName") """ for path in paths: - if path in openapi_schema.get("paths", {}) and "post" in openapi_schema["paths"][path]: - # Get the actual schema to extract ALL field definitions - schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref - actual_schema = openapi_schema.get("components", {}).get("schemas", {}).get(schema_name, {}) - schema_properties = actual_schema.get("properties", {}) - required_fields = actual_schema.get("required", []) + path_item = CustomOpenAPISpec._as_object( + CustomOpenAPISpec._as_object(openapi_schema.get("paths")).get(path) + ) + if "post" not in path_item: + continue - # Extract $defs and add them to components/schemas - # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI - if "$defs" in actual_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"]) + post_operation = CustomOpenAPISpec._as_object(path_item["post"]) - # Create an expanded inline schema instead of just a $ref - # This makes Swagger UI show all individual fields in the request body editor - expanded_schema = { - "type": "object", - "required": required_fields, - "properties": {}, - } + # Get the actual schema to extract ALL field definitions + schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref + components = CustomOpenAPISpec._as_object(openapi_schema.get("components")) + actual_schema = CustomOpenAPISpec._as_object( + CustomOpenAPISpec._as_object(components.get("schemas")).get(schema_name) + ) + schema_properties = CustomOpenAPISpec._as_object(actual_schema.get("properties")) + required_fields = actual_schema.get("required", []) - # Add all properties with their full definitions - for field_name, field_def in schema_properties.items(): - expanded_field = CustomOpenAPISpec._expand_field_definition(field_def) + # Extract $defs and add them to components/schemas + # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI + if "$defs" in actual_schema: + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, CustomOpenAPISpec._as_object(actual_schema["$defs"]) + ) - # Rewrite $defs references to use components/schemas instead - expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field) + # Create an expanded inline schema instead of just a $ref + # This makes Swagger UI show all individual fields in the request body editor + expanded_schema: JsonObject = { + "type": "object", + "required": required_fields, + "properties": { + field_name: CustomOpenAPISpec._expanded_request_field(field_name, field_def) + for field_name, field_def in schema_properties.items() + }, + } - # Add a simple example for the messages field - if field_name == "messages": - expanded_field["example"] = [{"role": "user", "content": "Hello, how are you?"}] + # Set the request body with the expanded schema + post_operation["requestBody"] = { + "required": True, + "content": {"application/json": {"schema": expanded_schema}}, + } - expanded_schema["properties"][field_name] = expanded_field - - # Set the request body with the expanded schema - openapi_schema["paths"][path]["post"]["requestBody"] = { - "required": True, - "content": {"application/json": {"schema": expanded_schema}}, - } - - # Keep any existing parameters (like path parameters) but remove conflicting query params - if "parameters" in openapi_schema["paths"][path]["post"]: - existing_params = openapi_schema["paths"][path]["post"]["parameters"] - # Only keep path parameters, remove query params that conflict with request body - filtered_params = [param for param in existing_params if param.get("in") == "path"] - openapi_schema["paths"][path]["post"]["parameters"] = filtered_params + # Keep any existing parameters (like path parameters) but remove conflicting query params + if "parameters" in post_operation: + # Only keep path parameters, remove query params that conflict with request body + post_operation["parameters"] = [ + param + for param in CustomOpenAPISpec._as_array(post_operation["parameters"]) + if CustomOpenAPISpec._as_object(param).get("in") == "path" + ] @staticmethod - def _move_defs_to_components(openapi_schema: dict[str, Any], defs: Mapping[str, Mapping[str, Any]]) -> None: + def _move_defs_to_components(openapi_schema: JsonObject, defs: Mapping[str, JsonValue]) -> None: """ Move $defs from Pydantic v2 schema to OpenAPI components/schemas. This makes the definitions resolvable in Swagger/OpenAPI viewers. @@ -146,23 +174,31 @@ class CustomOpenAPISpec: return # Ensure components/schemas exists - if "components" not in openapi_schema: - openapi_schema["components"] = {} - if "schemas" not in openapi_schema["components"]: - openapi_schema["components"]["schemas"] = {} + schemas: Final = CustomOpenAPISpec._components_schemas(openapi_schema) # Add each definition to components/schemas for def_name, def_schema in defs.items(): # Recursively rewrite any nested $defs references within this definition - rewritten_def = CustomOpenAPISpec._rewrite_defs_refs(def_schema) - openapi_schema["components"]["schemas"][def_name] = rewritten_def + schemas[def_name] = CustomOpenAPISpec._rewrite_defs_refs(def_schema) # If this definition also has $defs, process them recursively - if "$defs" in def_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"]) + def_object = CustomOpenAPISpec._as_object(def_schema) + if "$defs" in def_object: + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, CustomOpenAPISpec._as_object(def_object["$defs"]) + ) @staticmethod - def _rewrite_defs_refs(schema: Any) -> Any: + def _rewritten_defs_entry(key: str, value: JsonValue) -> JsonValue: + if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): + # Rewrite the reference to use components/schemas + def_name: Final = value.replace("#/$defs/", "") + return f"#/components/schemas/{def_name}" + # Recursively process nested structures + return CustomOpenAPISpec._rewrite_defs_refs(value) + + @staticmethod + def _rewrite_defs_refs(schema: JsonValue) -> JsonValue: """ Recursively rewrite $ref values from #/$defs/... to #/components/schemas/... This converts Pydantic v2 references to OpenAPI-compatible references. @@ -174,26 +210,17 @@ class CustomOpenAPISpec: Schema with rewritten references """ if isinstance(schema, dict): - result: Final = {} - for key, value in schema.items(): - if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): - # Rewrite the reference to use components/schemas - def_name = value.replace("#/$defs/", "") - result[key] = f"#/components/schemas/{def_name}" - elif key == "$defs": - # Remove $defs from the schema since they're moved to components - continue - else: - # Recursively process nested structures - result[key] = CustomOpenAPISpec._rewrite_defs_refs(value) - return result - elif isinstance(schema, list): + return { + key: CustomOpenAPISpec._rewritten_defs_entry(key, value) + for key, value in schema.items() + if key != "$defs" + } + if isinstance(schema, list): return [CustomOpenAPISpec._rewrite_defs_refs(item) for item in schema] - else: - return schema + return schema @staticmethod - def _extract_field_schema(field_def: dict[str, Any]) -> dict[str, Any]: + def _extract_field_schema(field_def: JsonObject) -> JsonValue: """ Extract a simple schema from a Pydantic field definition for parameter display. @@ -209,10 +236,10 @@ class CustomOpenAPISpec: # Handle anyOf (Optional fields in Pydantic v2) if "anyOf" in field_def: - any_of: Final = field_def["anyOf"] + any_of: Final = CustomOpenAPISpec._as_array(field_def["anyOf"]) # Find the non-null type for option in any_of: - if option.get("type") != "null": + if CustomOpenAPISpec._as_object(option).get("type") != "null": return option # Fallback to string if all else fails return {"type": "string"} @@ -221,7 +248,7 @@ class CustomOpenAPISpec: return {"type": "string"} @staticmethod - def _expand_field_definition(field_def: dict[str, object]) -> dict[str, object]: + def _expand_field_definition(field_def: JsonObject) -> JsonObject: """ Expand a Pydantic field definition for inline use in OpenAPI schema. This creates a full field definition that Swagger UI can render as individual form fields. @@ -237,12 +264,12 @@ class CustomOpenAPISpec: @staticmethod def add_request_schema( - openapi_schema: dict[str, object], + openapi_schema: JsonObject, model_class: type, schema_name: str, paths: Sequence[str], operation_name: str, - ) -> dict[str, object]: + ) -> JsonObject: """ Generic method to add a request schema to OpenAPI specification. @@ -282,8 +309,8 @@ class CustomOpenAPISpec: @staticmethod def add_chat_completion_request_schema( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -309,7 +336,7 @@ class CustomOpenAPISpec: return openapi_schema @staticmethod - def add_embedding_request_schema(openapi_schema: dict[str, object]) -> dict[str, object]: + def add_embedding_request_schema(openapi_schema: JsonObject) -> JsonObject: """ Add EmbeddingRequest schema to embedding endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -336,8 +363,8 @@ class CustomOpenAPISpec: @staticmethod def add_responses_api_request_schema( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add ResponsesAPIRequestParams schema to responses API endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -364,8 +391,8 @@ class CustomOpenAPISpec: @staticmethod def add_llm_api_request_schema_body( - openapi_schema: dict[str, object], - ) -> dict[str, object]: + openapi_schema: JsonObject, + ) -> JsonObject: """ Add LLM API request schema bodies to OpenAPI specification for documentation. @@ -376,12 +403,10 @@ class CustomOpenAPISpec: OpenAPI schema with added request body schemas """ # Add chat completion request schema - openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) + with_chat_completions: Final = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) # Add embedding request schema - openapi_schema = CustomOpenAPISpec.add_embedding_request_schema(openapi_schema) + with_embeddings: Final = CustomOpenAPISpec.add_embedding_request_schema(with_chat_completions) # Add responses API request schema - openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema(openapi_schema) - - return openapi_schema + return CustomOpenAPISpec.add_responses_api_request_schema(with_embeddings) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 3a1d18b48cc..554a6ae8d1a 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -6,9 +6,11 @@ import os import sys import tracemalloc from collections import Counter -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, NamedTuple, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query +from typing_extensions import ReadOnly from litellm import get_secret_str from litellm._logging import verbose_proxy_logger @@ -194,6 +196,42 @@ async def memory_usage_in_mem_cache_items( } +class _ProcessMemoryInfo(Protocol): + """The resident and virtual sizes psutil reports for a process.""" + + @property + def rss(self) -> int: ... + + @property + def vms(self) -> int: ... + + +class _ProcessHandle(Protocol): + """The psutil process handle members this module reads.""" + + def memory_info(self) -> _ProcessMemoryInfo: ... + + def memory_percent(self) -> float: ... + + +class _ProcessMemoryUsage(NamedTuple): + """Memory usage of a single worker process.""" + + resident_megabytes: float + virtual_megabytes: float + percent: float + + +def _process_memory_usage(process: _ProcessHandle) -> _ProcessMemoryUsage: + """Read resident/virtual megabytes and system memory share for ``process``.""" + memory_info: Final = process.memory_info() + return _ProcessMemoryUsage( + resident_megabytes=memory_info.rss / (1024 * 1024), + virtual_megabytes=memory_info.vms / (1024 * 1024), + percent=process.memory_percent(), + ) + + @router.get("/debug/memory/summary", include_in_schema=False) async def get_memory_summary( _: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -227,10 +265,9 @@ async def get_memory_summary( try: import psutil - process: Final = psutil.Process() - memory_info: Final = process.memory_info() - memory_mb: Final = memory_info.rss / (1024 * 1024) - memory_percent: Final = process.memory_percent() + usage: Final = _process_memory_usage(psutil.Process()) + memory_mb: Final = usage.resident_megabytes + memory_percent: Final = usage.percent process_memory = { "summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)", @@ -252,7 +289,7 @@ async def get_memory_summary( process_memory["error"] = str(e) # Get cache information - caches: Final[dict[str, Any]] = {} + caches: Final[dict[str, object]] = {} total_cache_items = 0 try: @@ -313,7 +350,7 @@ async def get_memory_summary( } -def _get_gc_statistics() -> dict[str, Any]: +def _get_gc_statistics() -> Mapping[str, object]: """Get garbage collector statistics.""" return { "enabled": gc.isenabled(), @@ -341,30 +378,42 @@ def _get_gc_statistics() -> dict[str, Any]: } -def _get_object_type_counts(top_n: int) -> tuple[int, list[dict[str, Any]]]: +class _ObjectTypeCount(TypedDict): + """One row of the tracked-object histogram.""" + + type: ReadOnly[str] + count: ReadOnly[int] + count_readable: ReadOnly[str] + + +def _type_name_counts(objects: Sequence[object]) -> Counter[str]: + """Count ``objects`` by the name of their type.""" + return Counter(type(obj).__name__ for obj in objects) + + +def _get_object_type_counts(top_n: int) -> tuple[int, list[_ObjectTypeCount]]: """Count objects by type and return total count and top N types.""" - type_counts: Final[Counter] = Counter() - total_objects = 0 + type_counts: Final = _type_name_counts(gc.get_objects()) - for obj in gc.get_objects(): - total_objects += 1 - obj_type = type(obj).__name__ - type_counts[obj_type] += 1 - - top_object_types: Final = [ + top_object_types: Final[list[_ObjectTypeCount]] = [ {"type": obj_type, "count": count, "count_readable": f"{count:,}"} for obj_type, count in type_counts.most_common(top_n) ] - return total_objects, top_object_types + return sum(type_counts.values()), top_object_types -def _get_uncollectable_objects_info() -> dict[str, Any]: +def _type_names(objects: Sequence[object]) -> Sequence[str]: + """The type name of each object in ``objects``.""" + return [type(obj).__name__ for obj in objects] + + +def _get_uncollectable_objects_info() -> Mapping[str, object]: """Get information about uncollectable objects (potential memory leaks).""" uncollectable: Final = gc.garbage return { "count": len(uncollectable), - "sample_types": [type(obj).__name__ for obj in uncollectable[:10]], + "sample_types": _type_names(uncollectable[:10]), "warning": ( "If count > 0, you may have reference cycles preventing garbage collection" if len(uncollectable) > 0 @@ -373,9 +422,11 @@ def _get_uncollectable_objects_info() -> dict[str, Any]: } -def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) -> dict[str, Any]: +def _get_cache_memory_stats( + user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache +) -> Mapping[str, object]: """Calculate memory usage for all caches.""" - cache_stats: Final[dict[str, Any]] = {} + cache_stats: Final[dict[str, object]] = {} try: # User API key cache user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) @@ -439,9 +490,9 @@ def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, r return cache_stats -def _get_router_memory_stats(llm_router) -> dict[str, Any]: +def _get_router_memory_stats(llm_router) -> Mapping[str, object]: """Get memory usage statistics for LiteLLM router.""" - litellm_router_memory: dict[str, Any] = {} + litellm_router_memory: dict[str, object] = {} try: if llm_router is not None: # Model list memory size @@ -505,7 +556,7 @@ def _get_router_memory_stats(llm_router) -> dict[str, Any]: return litellm_router_memory -def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> dict[str, Any] | None: +def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Mapping[str, object] | None: """Get process-level memory information using psutil.""" if not include_process_info: return None @@ -514,10 +565,10 @@ def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> dic import psutil process: Final = psutil.Process() - memory_info: Final = process.memory_info() - ram_usage_mb: Final = round(memory_info.rss / (1024 * 1024), 2) - virtual_memory_mb: Final = round(memory_info.vms / (1024 * 1024), 2) - memory_percent: Final = round(process.memory_percent(), 2) + usage: Final = _process_memory_usage(process) + ram_usage_mb: Final = round(usage.resident_megabytes, 2) + virtual_memory_mb: Final = round(usage.virtual_megabytes, 2) + memory_percent: Final = round(usage.percent, 2) return { "pid": worker_pid, diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 9fd24162f7e..213a697b3dd 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,13 +10,36 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast if TYPE_CHECKING: from litellm.router import Router +def configured_display_names( + entries: Sequence[tuple[str, str]], + llm_router: Router | None, +) -> Mapping[str, str]: + """response_id -> configured `model_info.display_name` for the listing entries + that have one. + + Metadata is looked up by each entry's internal lookup id (so team-scoped rows + resolve), while the returned map is keyed by the public response id the + Anthropic-shaped listing is built from. Entries without a configured name are + omitted so the listing falls back to the id itself. + """ + if llm_router is None: + return MappingProxyType({}) + resolved: Final = ( + (response_id, llm_router.get_configured_display_name(lookup_id)) for response_id, lookup_id in entries + ) + return MappingProxyType( + {response_id: display_name for response_id, display_name in resolved if display_name is not None} + ) + + class TeamModelNameTranslator: """Translates internal team routing keys to their public names for the model listing/retrieve responses. Stateless; the live router and general_settings diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 44b88aec47b..76982d30306 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Final, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload from pydantic import BaseModel @@ -9,6 +9,9 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +if TYPE_CHECKING: + from opentelemetry.trace import Span + T = TypeVar("T", bound=BaseModel) @@ -40,8 +43,8 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, *, model_type: type[T], @@ -51,8 +54,8 @@ class UserApiKeyCache(DualCache): @overload def get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: None = None, **kwargs: object, @@ -60,12 +63,12 @@ class UserApiKeyCache(DualCache): def get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: type[BaseModel] | None = None, **kwargs: object, - ) -> Any | BaseModel | None: + ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) cached: Final = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) @@ -86,8 +89,8 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, *, model_type: type[T], @@ -97,8 +100,8 @@ class UserApiKeyCache(DualCache): @overload async def async_get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: None = None, **kwargs: object, @@ -106,12 +109,12 @@ class UserApiKeyCache(DualCache): async def async_get_cache( self, - key: object, - parent_otel_span: object = None, + key: str, + parent_otel_span: Span | None = None, local_only: bool = False, model_type: type[BaseModel] | None = None, **kwargs: object, - ) -> Any | BaseModel | None: + ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) cached: Final = await super().async_get_cache( @@ -131,12 +134,12 @@ class UserApiKeyCache(DualCache): return None return decoded - def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): + def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) - async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object): + async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 202a95ba29b..e6880d521f1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -211,7 +211,7 @@ class DBSpendUpdateWriter: org_id: str | None, # Completion object fields kwargs: dict | None, - completion_response: litellm.ModelResponse | Any | Exception | None, + completion_response: object, start_time: datetime | None, end_time: datetime | None, response_cost: float | None, @@ -323,7 +323,7 @@ class DBSpendUpdateWriter: async def _enqueue_tool_usage_transaction( self, payload: SpendLogsPayload, - completion_response: "litellm.ModelResponse | Any | Exception | None", + completion_response: object, prisma_client: "PrismaClient | None", kwargs: "dict | None" = None, ) -> None: @@ -396,7 +396,7 @@ class DBSpendUpdateWriter: def _enqueue_tool_registry_upsert( self, kwargs: dict | None, - completion_response: Any | None, + completion_response: object, hashed_token: str | None = None, team_id: str | None = None, ) -> None: @@ -849,7 +849,7 @@ class DBSpendUpdateWriter: return # Parse tags from JSON string - tags = [] + tags: Sequence[object] = [] if isinstance(request_tags, str): tags = safe_json_loads(request_tags, default=[]) if not tags: @@ -2260,7 +2260,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.") return - request_tags = [] + request_tags: Sequence[str] = [] if isinstance(payload["request_tags"], str): request_tags = json.loads(payload["request_tags"]) elif isinstance(payload["request_tags"], list): diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 1a39016b3a3..01f66e4f3c5 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -82,10 +82,30 @@ CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset( "pool_timeout", "connect_timeout", "socket_timeout", + "max_idle_connection_lifetime", "pgbouncer", } ) +# Quaint never tests pooled connections on checkout and keeps them idle for +# 300s by default, past many infra idle timeouts, so dead sockets surface as +# `Error { kind: Closed }`. 60s recycles them first; explicit values win. +DEFAULT_MAX_IDLE_CONNECTION_LIFETIME: Final = 60 +IDLE_LIFETIME_DEFAULT_PARAMS: Final[Mapping[str, int]] = MappingProxyType( + {"max_idle_connection_lifetime": DEFAULT_MAX_IDLE_CONNECTION_LIFETIME} +) + + +def idle_lifetime_params(configured: float | None) -> Mapping[str, str | int | float]: + """The `max_idle_connection_lifetime` to add to URLs that do not pin one. + + Applied via ``add_missing_query_params`` so a URL-pinned value always wins, + whether the operator configured `database_max_idle_connection_lifetime` or not. + """ + if configured is None: + return IDLE_LIFETIME_DEFAULT_PARAMS + return MappingProxyType({"max_idle_connection_lifetime": configured}) + def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) -> str: """Return ``url`` with the ``params`` it does not already carry appended. diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py index 3716d00774f..2c27531cea1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py @@ -162,10 +162,10 @@ class AktoGuardrail(CustomGuardrail): def build_request_body( inputs: GenericGuardrailAPIInputs, request_data: dict | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the LLM request body from guardrail inputs (messages, model, tools).""" model: Final = inputs.get("model", "") or "" - body: Final[dict[str, Any]] = {"model": model} + body: Final[dict[str, object]] = {"model": model} structured: Final = inputs.get("structured_messages") if structured: @@ -194,7 +194,7 @@ class AktoGuardrail(CustomGuardrail): def build_response_body( inputs: GenericGuardrailAPIInputs, request_data: dict | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the LLM response body, preferring the actual model response if available.""" model_response: Final = request_data.get("response") if request_data else None if model_response is not None and hasattr(model_response, "model_dump"): @@ -224,7 +224,7 @@ class AktoGuardrail(CustomGuardrail): *, status_code: int = 200, include_response: bool = False, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build the flat MIRRORING payload sent to Akto's HTTP proxy endpoint. All body fields use double-encoding: json.dumps({"body": json.dumps(actual_body)}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/alice/__init__.py new file mode 100644 index 00000000000..75ea16f7a88 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/__init__.py @@ -0,0 +1,34 @@ +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .alice import AliceGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _alice_guardrail_callback: Final = AliceGuardrail( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"), + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(_alice_guardrail_callback) + return _alice_guardrail_callback + + +guardrail_initializer_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.ALICE.value: initialize_guardrail, +} + + +guardrail_class_registry: Final = { # mutable-ok: module-level registry, built once and never mutated + SupportedGuardrailIntegrations.ALICE.value: AliceGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py new file mode 100644 index 00000000000..27018769909 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py @@ -0,0 +1,369 @@ +# +-------------------------------------------------------------+ +# +# Use Alice for your LLM calls +# https://alice.io/ +# +# +-------------------------------------------------------------+ + +import json +import os +from collections.abc import Mapping +from typing import ( + TYPE_CHECKING, + Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml + Final, + Literal, + Optional, +) + +import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException, Timeout +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME: Final = "alice" + +_DEFAULT_API_BASE: Final = "https://api.alice.io" +_EVALUATE_PATH: Final = "/v2/evaluate/litellm" + +_VERDICT_ALLOW: Final = "ALLOW" +_VERDICT_BLOCK: Final = "BLOCK" +_VERDICT_MASK: Final = "MASK" +_VERDICT_DETECT: Final = "DETECT" +_KNOWN_VERDICTS: Final = frozenset({_VERDICT_ALLOW, _VERDICT_BLOCK, _VERDICT_MASK, _VERDICT_DETECT}) + +_DEFAULT_BLOCK_MESSAGE: Final = "Blocked by your organization's content policy." + +# apply_guardrail selects nothing: it forwards whichever of these came populated and lets Alice +# decide what is worth evaluating. Only skip the call when every one of them is empty — there is +# then genuinely nothing to send. +_SELECTABLE_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls", "structured_messages") + +# Caps on the outbound copy of request_data. A payload deeper or wider than this is malformed +# rather than large, and serializing it would cost more than the evaluation it feeds. +_MAX_DEPTH: Final = 12 +_MAX_ITEMS: Final = 5000 + +# request_data carries the caller's raw credentials under these keys, at any nesting depth — +# a real captured payload puts inbound headers at request_data["proxy_server_request"]["headers"], +# again under ["metadata"]["headers"] / ["litellm_metadata"]["headers"], and again under +# ["metadata"]["requester_metadata"]["headers"], any of which can carry an Authorization or +# x-api-key value. LiteLLM's own spend-log sanitizer excludes `secret_fields` for the same reason +# (spend_tracking_utils._SENSITIVE_REQUEST_BODY_KEYS): `secret_fields.raw_headers` holds the +# caller's Authorization / x-api-key in the clear, and `api_key` can carry a forwarded provider +# credential. Stripping by key name rather than by path means a new nesting path can never +# reintroduce the leak. Posting any of these to a third-party guardrail endpoint would be worse +# than what the proxy already refuses to persist in its own audit trail — so none of them leave +# the process. +_CREDENTIAL_KEYS_TO_STRIP: Final = frozenset( + {"secret_fields", "api_key", "raw_headers", "headers", "provider_specific_header"} +) + + +class AliceReplacement(TypedDict): + """A masked substitution, positional against the texts that were submitted.""" + + index: ReadOnly[NotRequired[int]] + text: ReadOnly[NotRequired[str]] + + +class AliceVerdict(TypedDict): + """Body returned by Alice's LiteLLM evaluate endpoint.""" + + verdict: ReadOnly[NotRequired[str]] + categories: ReadOnly[NotRequired["tuple[str, ...]"]] + correlation_id: ReadOnly[NotRequired[str]] + message: ReadOnly[NotRequired[str]] + replacements: ReadOnly[NotRequired["tuple[AliceReplacement, ...]"]] + + +class AliceGuardrailMissingSecrets(Exception): + """Raised when the Alice API key is not configured.""" + + +class AliceGuardrail(CustomGuardrail): + """ + Alice — policy-based guardrails for prompts and model responses. + + This forwards the hook's arguments as it received them and enforces the verdict that comes + back, with one deliberate exception: any key named `secret_fields`, `api_key`, `raw_headers`, + `headers`, or `provider_specific_header` is dropped from `request_data` at any nesting depth + before it is serialized, and never reaches Alice. Short of that, it selects nothing and + renames nothing: which parts of a conversation are worth evaluating, and how a verdict is + reached, are decided by Alice — so changing either is a change on their side rather than a + LiteLLM upgrade. A batch with nothing selectable at all (no `texts`, `images`, `tools`, + `tool_calls`, or `structured_messages`) still skips the call, since there would be nothing to + send. + + Known limitation: the unified guardrail's `streaming_transform_mode` defaults to + `block_only`, whose streaming path discards any returned text rewrite. A MASK verdict is + therefore a no-op on a streamed response — the original, unmasked text still reaches the + caller — while BLOCK continues to function on both streamed and non-streamed responses. + This is `during_call`'s documented behavior generally, not specific to Alice; configure a + masking-aware `streaming_transform_mode` if that gap matters for your traffic. + + Alice evaluates against policies configured per *application*, and one proxy typically fronts + several, so the application is named on the virtual key rather than in this config: + + curl $PROXY/key/generate -H "Authorization: Bearer $LITELLM_MASTER_KEY" \\ + -d '{"key_alias": "payments-bot", + "metadata": {"alice_app_id": "payments-bot"}}' + + Alice reads that off the authenticated key. Because the proxy strips caller-supplied + `user_api_key_*` from the request before a guardrail sees it, a caller cannot point its own + traffic at an application with laxer policies than the one its key was issued for. + + Configuration example (litellm config YAML): + guardrails: + - guardrail_name: alice + litellm_params: + guardrail: alice + mode: [pre_call, post_call] + api_key: os.environ/ALICE_API_KEY + api_base: https://api.alice.io # optional + unreachable_fallback: fail_closed # optional + """ + + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + **kwargs: Any, # kwargs-ok: forwarded verbatim to CustomGuardrail.__init__, whose param list is wide and evolving + ) -> None: + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + + alice_api_key: Final = api_key or os.environ.get("ALICE_API_KEY") + if not alice_api_key: + raise AliceGuardrailMissingSecrets( + "Alice API key is required. Set the `ALICE_API_KEY` environment variable or " + "pass `api_key` in the guardrail config." + ) + self.alice_api_key: str = alice_api_key + + base: Final = (api_base or os.environ.get("ALICE_API_BASE") or _DEFAULT_API_BASE).rstrip("/") + self.api_base: str = f"{base}{_EVALUATE_PATH}" + self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback + + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ # mutable-ok: CustomGuardrail.__init__ requires a list here + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + + super().__init__(**kwargs) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: overrides CustomGuardrail.apply_guardrail's plain-dict contract + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + if not any(inputs.get(field) for field in _SELECTABLE_INPUT_FIELDS): + return inputs + + try: + verdict: AliceVerdict = await self._evaluate( + inputs=inputs, request_data=request_data, input_type=input_type + ) + except Timeout as e: + return self._on_unreachable(e, inputs) + except httpx.HTTPStatusError as e: + status_code: Final = getattr(getattr(e, "response", None), "status_code", None) + # Any 5xx is an outage on Alice's side, not our misconfiguration — route the whole + # class through the configured policy. A 4xx (rejected credential, bad request) is + # ours to fix and must never fail open, so it is deliberately left to propagate. + if isinstance(status_code, int) and 500 <= status_code < 600: + return self._on_unreachable(e, inputs) + raise + except httpx.RequestError as e: + return self._on_unreachable(e, inputs) + except (json.JSONDecodeError, UnicodeDecodeError, TypeError) as e: + # A body that cannot be decoded, cannot be parsed as JSON, or parses to something + # other than an object, is as unreachable as a dropped connection: this deployment's + # policy decides, not a raw exception. UnicodeDecodeError is named explicitly because + # it is a sibling of JSONDecodeError under ValueError, not a subclass of it. + return self._on_unreachable(e, inputs) + + return self._enforce(verdict, inputs) + + async def _evaluate( + self, + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: str, + ) -> AliceVerdict: + response: Final = await self.async_handler.post( + url=self.api_base, + json={ # mutable-ok: one-shot HTTP request body, never mutated after construction + "input_type": input_type, + "inputs": _json_safe(inputs), + "request_data": _json_safe(request_data, strip_keys=_CREDENTIAL_KEYS_TO_STRIP), + }, + headers={ # mutable-ok: one-shot HTTP headers, never mutated after construction + "Content-Type": "application/json", + "af-api-key": self.alice_api_key, + }, + ) + response.raise_for_status() + body = response.json() + if not isinstance(body, dict): + raise TypeError("Alice returned a non-object body") + return body + + def _enforce(self, verdict: AliceVerdict, inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs: + """Act on the verdict. An answer we cannot read is treated as unavailable, never as a pass.""" + name: Final = verdict.get("verdict") + if name not in _KNOWN_VERDICTS: + return self._on_unreachable(ValueError(f"unrecognized verdict: {name!r}"), inputs) + + if name == _VERDICT_BLOCK: + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message=verdict.get("message") or _DEFAULT_BLOCK_MESSAGE, + should_wrap_with_default_message=False, + blocked_content=True, + ) + + if name == _VERDICT_DETECT: + # Recorded by Alice and allowed through. The correlation id is what ties this request + # to that record; the evaluated text itself is never logged. + verbose_proxy_logger.warning( + "Alice guardrail: detection recorded, request allowed (correlation_id=%s, categories=%s)", + verdict.get("correlation_id"), + verdict.get("categories"), + ) + return inputs + + if name == _VERDICT_MASK: + self._apply_replacements(verdict, inputs) + + return inputs + + def _apply_replacements(self, verdict: AliceVerdict, inputs: GenericGuardrailAPIInputs) -> None: + """ + Write each replacement onto the text it names. + + Only `texts` is touched. The chat translation layer maps a returned `texts` list back onto + the request positionally, but takes a different branch entirely when `structured_messages` + comes back as a new object — which would drop these edits. + + All-or-nothing: a single out-of-range or malformed replacement blocks the whole verdict + rather than being silently skipped, so content Alice meant to replace can never reach the + model unmasked alongside content that was replaced. + """ + texts: Final = inputs.get("texts") or [] # mutable-ok: empty-list fallback, replaced wholesale below + replacements: Final = verdict.get("replacements") or [] # mutable-ok: empty-list fallback for iteration only + + if not replacements: + raise self._mask_rejected(verdict) + + for replacement in replacements: + index = replacement.get("index") + text = replacement.get("text") + if not (isinstance(index, int) and isinstance(text, str) and 0 <= index < len(texts)): + raise self._mask_rejected(verdict) + texts[index] = text # mutable-ok: item assignment into the local working copy above + + inputs["texts"] = texts + + def _mask_rejected(self, verdict: AliceVerdict) -> GuardrailRaisedException: + """A MASK verdict that cannot be applied in full is refused outright, never partially — + see `_apply_replacements`.""" + return GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message=verdict.get("message") or _DEFAULT_BLOCK_MESSAGE, + should_wrap_with_default_message=False, + blocked_content=True, + ) + + def _on_unreachable(self, error: Exception, inputs: GenericGuardrailAPIInputs) -> GenericGuardrailAPIInputs: + """Apply the configured policy when Alice cannot be reached or cannot be understood.""" + if self.unreachable_fallback == "fail_open": + verbose_proxy_logger.critical( + "Alice guardrail unreachable, allowing request per unreachable_fallback: %s", + error, + ) + return inputs + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message="Alice guardrail is unavailable and this request cannot be checked", + should_wrap_with_default_message=False, + ) from error + + @staticmethod + def get_config_model() -> type | None: + from litellm.types.proxy.guardrails.guardrail_hooks.alice import ( + AliceGuardrailConfigModel, + ) + + return AliceGuardrailConfigModel + + +def _json_safe( + value: object, + depth: int = 0, + seen: frozenset[int] = frozenset(), + strip_keys: frozenset[str] = frozenset(), +) -> object: + """ + Copy `value` into something `json.dumps` accepts, dropping only what cannot cross. + + `request_data` carries live Python objects — an OpenTelemetry span among them — so it cannot + be serialized as it stands. What is dropped is decided by a mechanical rule rather than a + field list: a list drifts from what the far side needs, a rule cannot. Serializing naively + raises, and that error would be read as "guardrail unavailable" on every single request. + + `strip_keys` drops a dict key by name at every depth it appears, not just the root — a caller + passes `_CREDENTIAL_KEYS_TO_STRIP` here so a credential nested under any path is caught the + same way a top-level one is, without maintaining a list of paths. The source object is never + mutated: every branch below builds a new container. + """ + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if depth >= _MAX_DEPTH or id(value) in seen: + return None + + nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately + + if isinstance(value, dict): + out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is + for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view + if isinstance(key, str) and key not in strip_keys: + out[key] = _json_safe(item, depth + 1, nested, strip_keys) + return out + + if isinstance(value, (list, tuple, set, frozenset)): + return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use + _json_safe(item, depth + 1, nested, strip_keys) + for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view + ] + + dump: Final = getattr(value, "model_dump", None) + if callable(dump): + try: + return _json_safe(dump(mode="json"), depth + 1, nested, strip_keys) + except Exception: # noqa: BLE001 # a model that will not dump is one we drop + return None + + # Everything json.dumps handles natively — str, int, float, bool, None, dict, list — is + # caught above, and a dict/list subclass is caught by isinstance. So whatever reaches here + # (bytes, datetime, an OpenTelemetry span) cannot cross the wire. + return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index a6635ea0776..30526d30dc5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -31,7 +31,11 @@ from litellm.caching import DualCache from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.litellm_logging import ( + _get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -212,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]: return redacted if isinstance(redacted, list) else assessments +_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) + + +def _is_responses_api_route(request_route: str | None) -> bool: + if request_route is None: + return False + call_types: Final = get_call_types_for_route(request_route) + return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types) + + class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL. @@ -252,7 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` # routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail. - self.checks: dict[str, Any] | None = self._normalize_checks(checks) + self.checks: dict[str, object] | None = self._normalize_checks(checks) # Per-check block thresholds; a score >= threshold blocks. None => the # check is detect-only (logged, never blocks). self.content_filter_threshold = content_filter_threshold @@ -321,7 +335,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ] @staticmethod - def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, Any] | None: + def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, object] | None: """Normalize the configured `checks` into a plain dict for the API body. Accepts a pydantic ``BedrockChecksConfigModel`` or a raw dict; drops None / @@ -372,7 +386,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _create_bedrock_output_content_request( self, - response: Any | ModelResponse, + response: object, messages: list[AllMessageValues] | None = None, ) -> BedrockRequest: """ @@ -396,9 +410,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_request["content"] = bedrock_request_content return bedrock_request - def _build_response_content_items( - self, response: Any | ModelResponse, has_grounding: bool - ) -> list[BedrockContentItem]: + def _build_response_content_items(self, response: object, has_grounding: bool) -> list[BedrockContentItem]: """Build content item(s) from the model response. When the request supplied grounding, the response is qualified ``guard_content`` so Bedrock can score it. """ @@ -422,7 +434,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self, source: Literal["INPUT", "OUTPUT"], messages: list[AllMessageValues] | None = None, - response: Any | ModelResponse | None = None, + response: object | None = None, ) -> BedrockRequest: """ Convert the litellm messages/response to the bedrock request format. @@ -945,7 +957,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _apply_guardrail_content_with_chunking( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1083,7 +1095,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _post_apply_guardrail_content_with_retry( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1133,7 +1145,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _post_apply_guardrail_content( self, content: Sequence[BedrockContentItem], - base_request_data: Mapping[str, Any], + base_request_data: Mapping[str, object], credentials: "Credentials", aws_region_name: str, api_key: str | None, @@ -1172,11 +1184,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): aws_region_name=aws_region_name, api_key=api_key, ) + headers_dict: Final = dict(prepared_request.headers) # mutable-ok: the masking helper requires a dict verbose_proxy_logger.debug( "Bedrock AI request body: %s, url %s, headers: %s", bedrock_request_data, prepared_request.url, - prepared_request.headers, + _get_masked_values(headers_dict), ) httpx_response: Final = await self._sign_and_post( @@ -1861,7 +1874,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() credentials, aws_region_name = self._load_credentials() - body: Final[dict[str, Any]] = {"messages": checks_messages, "checks": self.checks} + body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} api_key: Final[str | None] = request_data.get("api_key") if request_data else None prepared_request: Final = self._prepare_request( @@ -2343,7 +2356,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): guardrail_name=self.guardrail_name, ) - detail: Final[dict[str, Any]] = { + detail: Final[dict[str, object]] = { "error": "Violated guardrail policy", "bedrock_guardrail_response": bedrock_guardrail_output_text, } @@ -2707,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): yield streamed_chunk return + # Responses-API events are neither chat-completions chunks nor raw + # Anthropic SSE, so the assembly below cannot scan them; the unified + # guardrail's translation layer can, with buffering semantics kept. + if _is_responses_api_route(user_api_key_dict.request_route): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=True, + ): + yield translated_chunk + return + # Import here to avoid circular imports from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder @@ -2902,7 +2933,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return updated_messages def _mask_content_list( - self, content_list: list[Any], masked_texts: list[str], masking_index: int + self, content_list: Sequence[object], masked_texts: list[str], masking_index: int ) -> tuple[list[Any], int]: """ Apply masking to a list of content items. @@ -2915,7 +2946,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Returns: Updated content list with masked items """ - new_content: Final[list[dict | str]] = [] + new_content: Final[list[dict[str, object] | str]] = [] for item in content_list: if isinstance(item, dict) and "text" in item: new_item = item.copy() @@ -2934,7 +2965,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _apply_masking_to_response( self, - response: ModelResponse | Any, + response: object, bedrock_guardrail_response: BedrockGuardrailResponse, ) -> None: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index 8398ec9f141..5a6be1089b6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -5,8 +5,9 @@ The public guardrail class imports this private mixin from while preserving the existing public import path. """ +from collections.abc import Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -23,7 +24,7 @@ if TYPE_CHECKING: from .cisco_ai_defense import _ScanContext -def _serialize_mcp_content_item(item: object) -> dict[str, Any]: +def _serialize_mcp_content_item(item: object) -> dict[str, object]: """Serialize an MCP content item to a JSON-friendly dict. Handles raw dicts, MCP SDK Pydantic models, and simple ``.text`` objects. @@ -57,7 +58,7 @@ class _CiscoAIDefenseMcpMixin: def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool: ... - async def _post_inspection(self, url: str, payload: dict[str, Any], surface: str) -> dict[str, Any]: ... + async def _post_inspection(self, url: str, payload: dict[str, object], surface: str) -> dict[str, object]: ... def _handle_api_error( self, @@ -67,16 +68,16 @@ class _CiscoAIDefenseMcpMixin: start_time: datetime | None = ..., surface: str = ..., direction: str = ..., - ) -> dict[str, Any]: ... + ) -> dict[str, object]: ... def _finalize_inspection( self, - inspect_response: dict[str, Any], + inspect_response: dict[str, object], request_data: dict, context: "_ScanContext", start_time: datetime, response_obj: object = ..., - ) -> dict[str, Any]: ... + ) -> dict[str, object]: ... # ------------------------------------------------------------------ # MCP post-tool hook (dispatcher contract) @@ -95,7 +96,7 @@ class _CiscoAIDefenseMcpMixin: if self.inspection_type != "mcp": return None - request_data: Final[dict[str, Any]] = {} + request_data: Final[dict[str, object]] = {} for key in ( "name", "litellm_call_id", @@ -188,9 +189,9 @@ class _CiscoAIDefenseMcpMixin: original_hidden: Final = getattr(original_response_obj, "hidden_params", None) if isinstance(original_hidden, HiddenParams): - hidden_params: Any = original_hidden + hidden_params: HiddenParams = original_hidden else: - response_cost: Final = getattr(original_hidden, "response_cost", None) + response_cost: Final[float | None] = getattr(original_hidden, "response_cost", None) hidden_params = HiddenParams(response_cost=response_cost) if response_cost is not None else HiddenParams() return MCPPostCallResponseObject( @@ -200,11 +201,11 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _replace_mcp_tool_response(response_obj: object, replacement_obj: object) -> bool: - replacement: Final = getattr(replacement_obj, "mcp_tool_call_response", None) + replacement: Final[list[object] | None] = getattr(replacement_obj, "mcp_tool_call_response", None) if replacement is None: return False - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: if _CiscoAIDefenseMcpMixin._replace_mcp_tool_response(inner, replacement_obj): return True @@ -276,7 +277,7 @@ class _CiscoAIDefenseMcpMixin: self, data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: del user_api_key_dict # carried via logging metadata, not the wire payload url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_mcp_request_payload(data=data) @@ -312,7 +313,7 @@ class _CiscoAIDefenseMcpMixin: response: object, user_api_key_dict: UserAPIKeyAuth | None = None, redact_response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: del user_api_key_dict # carried via logging metadata, not the wire payload url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_mcp_response_payload( @@ -349,7 +350,7 @@ class _CiscoAIDefenseMcpMixin: def _build_mcp_request_payload( self, data: dict, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """Build the JSON-RPC ``tools/call`` envelope sent to ``/inspect/mcp``. The Cisco AI Defense MCP inspect endpoint expects the JSON-RPC @@ -390,7 +391,7 @@ class _CiscoAIDefenseMcpMixin: self, request_data: dict, response: object, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """Build the MCP response-inspection body sent to ``/inspect/mcp``.""" request_payload: Final = self._build_mcp_request_payload(data=request_data) if request_payload is None: @@ -415,7 +416,7 @@ class _CiscoAIDefenseMcpMixin: return payload @staticmethod - def _hydrate_mcp_tool_context(request_data: dict[str, Any]) -> None: + def _hydrate_mcp_tool_context(request_data: dict[str, object]) -> None: metadata = request_data.get("mcp_tool_call_metadata") if metadata is None: nested: Final = request_data.get("metadata") or request_data.get("litellm_metadata") @@ -440,7 +441,7 @@ class _CiscoAIDefenseMcpMixin: request_data.setdefault("server_name", server_name) @staticmethod - def _normalize_mcp_response(response: object) -> dict[str, Any] | None: + def _normalize_mcp_response(response: object) -> dict[str, object] | None: """Normalize an MCP tool response into a JSON-RPC envelope. Handles JSON-RPC dicts, raw content lists, MCP SDK models, and @@ -502,10 +503,10 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _build_mcp_result( - content: list[Any], + content: Sequence[object], source: object = None, - ) -> dict[str, Any]: - result: Final[dict[str, Any]] = {"content": [_serialize_mcp_content_item(item) for item in content]} + ) -> dict[str, object]: + result: Final[dict[str, object]] = {"content": [_serialize_mcp_content_item(item) for item in content]} for key in ("structuredContent", "isError"): value = source.get(key) if isinstance(source, dict) else getattr(source, key, None) if value is not None and (key != "isError" or isinstance(value, bool)): @@ -522,7 +523,7 @@ class _CiscoAIDefenseMcpMixin: if response_obj is None: return False - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: return _CiscoAIDefenseMcpMixin._set_mcp_tool_response_text(inner, text) @@ -559,7 +560,7 @@ class _CiscoAIDefenseMcpMixin: pass elif isinstance(response_obj, dict): result: Final = response_obj.get("result") - target: Final[dict[Any, Any]] = result if isinstance(result, dict) else response_obj + target: Final[dict[object, object]] = result if isinstance(result, dict) else response_obj if "structuredContent" in target: target["structuredContent"] = replacement replaced = True @@ -567,11 +568,11 @@ class _CiscoAIDefenseMcpMixin: return replaced @staticmethod - def _coerce_to_content_list(response_obj: object) -> list[Any] | None: + def _coerce_to_content_list(response_obj: object) -> list[object] | None: """Find the MCP content list inside supported response shapes.""" if response_obj is None: return None - inner: Final = getattr(response_obj, "mcp_tool_call_response", None) + inner: Final[object | None] = getattr(response_obj, "mcp_tool_call_response", None) if inner is not None: return _CiscoAIDefenseMcpMixin._coerce_to_content_list(inner) content: Final = getattr(response_obj, "content", None) @@ -594,8 +595,8 @@ class _CiscoAIDefenseMcpMixin: @staticmethod def _extract_sanitized_mcp_arguments( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: dict[str, object], + ) -> dict[str, object] | None: """Pull sanitized MCP tool-call arguments off the verdict. Cisco can return them at the top level (``params.arguments``) or diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 955a868a0d6..48832f8ed5e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -2,9 +2,10 @@ import os import time -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -27,6 +28,33 @@ if TYPE_CHECKING: GRAYSWAN_BLOCK_ERROR_MSG: Final = "Blocked by Gray Swan Guardrail" +class _GraySwanMonitorResponse(TypedDict): + """Body returned by Gray Swan's `/cygnal/monitor` endpoint.""" + + violation: ReadOnly[NotRequired[float | None]] + violated_rules: ReadOnly[NotRequired[list[object]]] + violated_rule_descriptions: ReadOnly[NotRequired[list[object]]] + mutation: ReadOnly[NotRequired[bool | None]] + ipi: ReadOnly[NotRequired[bool | None]] + + +class _GraySwanMonitorHTTPResponse(Protocol): + def raise_for_status(self) -> object: ... + + def json(self) -> _GraySwanMonitorResponse: ... + + +class _GraySwanMonitorHTTPClient(Protocol): + async def post( + self, + *, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> _GraySwanMonitorHTTPResponse: ... + + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -77,7 +105,9 @@ class GraySwanGuardrail(CustomGuardrail): guardrail_timeout: float | None = 30.0, **kwargs: Any, ) -> None: - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler: _GraySwanMonitorHTTPClient = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) api_key_value: Final = api_key or os.getenv("GRAYSWAN_API_KEY") if not api_key_value: @@ -266,7 +296,7 @@ class GraySwanGuardrail(CustomGuardrail): # Legacy Test Interface (for backward compatibility) # ------------------------------------------------------------------ - async def run_grayswan_guardrail(self, payload: dict) -> dict[str, Any]: + async def run_grayswan_guardrail(self, payload: dict[str, object]) -> _GraySwanMonitorResponse: """ Run the GraySwan guardrail on a payload. @@ -285,7 +315,7 @@ class GraySwanGuardrail(CustomGuardrail): def _process_grayswan_response( self, - response_json: dict, + response_json: _GraySwanMonitorResponse, data: dict | None = None, hook_type: GuardrailEventHooks | None = None, ) -> None: @@ -385,7 +415,7 @@ class GraySwanGuardrail(CustomGuardrail): # Core GraySwan API interaction # ------------------------------------------------------------------ - async def _call_grayswan_api(self, payload: dict) -> dict[str, Any]: + async def _call_grayswan_api(self, payload: dict[str, object]) -> _GraySwanMonitorResponse: """Call the GraySwan monitoring API.""" headers: Final = self._prepare_headers() @@ -406,7 +436,7 @@ class GraySwanGuardrail(CustomGuardrail): def _process_response_internal( self, - response_json: dict[str, Any], + response_json: _GraySwanMonitorResponse, request_data: dict, inputs: GenericGuardrailAPIInputs, is_output: bool, @@ -534,8 +564,8 @@ class GraySwanGuardrail(CustomGuardrail): dynamic_body: dict, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"] = None, - ) -> dict[str, Any] | None: - payload: Final[dict[str, Any]] = {"messages": messages} + ) -> dict[str, object] | None: + payload: Final[dict[str, object]] = {"messages": messages} categories: Final = dynamic_body.get("categories") or self.categories if categories: @@ -563,13 +593,13 @@ class GraySwanGuardrail(CustomGuardrail): {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers ) if cleaned_litellm_metadata: - sanitized: Final = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) if isinstance(sanitized, dict) and sanitized: payload["litellm_metadata"] = sanitized return payload - def _format_violation_message(self, detection_info: Any, is_output: bool = False) -> str: + def _format_violation_message(self, detection_info: object, is_output: bool = False) -> str: """ Format detection info into a user-friendly violation message. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index ea022510309..cf5da27e9ca 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -8,6 +8,7 @@ import json import os import uuid +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict try: @@ -128,7 +129,7 @@ class LassoGuardrail(CustomGuardrail): @staticmethod def _extract_tool_call_fields( - call: Any, + call: object, ) -> tuple[str | None, str | None, dict[str, object] | None]: """Extract (call_id, name, parsed_input) from a tool call. @@ -476,7 +477,7 @@ class LassoGuardrail(CustomGuardrail): def _map_masked_messages_back( self, original_messages: list[dict[str, Any]], - masked_messages: list[dict[str, Any]], + masked_messages: Sequence[Mapping[str, object]], ) -> list[dict[str, object]]: """Map Lasso-format masked messages back onto the original OpenAI-format messages. @@ -638,7 +639,7 @@ class LassoGuardrail(CustomGuardrail): }, ) - def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _expand_messages_for_classification(self, messages: list[dict[str, Any]]) -> list[dict[str, object]]: """ Convert raw OpenAI-format messages to Lasso API format with content blocks. @@ -646,7 +647,7 @@ class LassoGuardrail(CustomGuardrail): - role=tool messages → developer role + tool_result block - plain text messages pass through unchanged """ - expanded: Final[list[dict[str, Any]]] = [] + expanded: Final[list[dict[str, object]]] = [] for msg in messages: role = msg.get("role", "") content = msg.get("content") @@ -917,7 +918,7 @@ class LassoGuardrail(CustomGuardrail): def _apply_masking_to_model_response( self, model_response: litellm.ModelResponse, - masked_messages: list[dict[str, Any]], + masked_messages: Sequence[Mapping[str, object]], ) -> None: """Apply masking to the actual model response when mask=True and masked content is available.""" # Index masked tool_use blocks by id for O(1) lookup. diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index e2d7c06f7c5..a269ad31a6b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -80,7 +80,7 @@ import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey -from typing_extensions import NotRequired, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -89,6 +89,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral @@ -107,6 +108,19 @@ class _JWTDecodeKwargs(TypedDict): issuer: NotRequired[str] +class _DebugHeaderClaims(TypedDict, total=False): + sub: ReadOnly[object] + iss: ReadOnly[object] + exp: ReadOnly[object] + scope: ReadOnly[str] + + +class _SignedClaimSummary(TypedDict): + sub: ReadOnly[object] + act: ReadOnly[Mapping[str, object]] + exp: ReadOnly[object] + + # Module-level singleton for the JWKS discovery endpoint to access. _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None @@ -265,7 +279,8 @@ class MCPJWTSigner(CustomGuardrail): **kwargs: Any, ) -> None: kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - super().__init__(**kwargs) + base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs + super().__init__(**base_kwargs) # --- Signing key setup --- key_material: Final = os.environ.get(self.SIGNING_KEY_ENV) @@ -677,7 +692,7 @@ class MCPJWTSigner(CustomGuardrail): data: dict, jwt_claims: Mapping[str, object] | None = None, call_type: CallTypesLiteral | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Build JWT claims for the outbound MCP access token. @@ -752,7 +767,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ @staticmethod - def _build_debug_header(claims: dict[str, Any], kid: str) -> str: + def _build_debug_header(claims: _DebugHeaderClaims, kid: str) -> str: """ Build the x-litellm-mcp-debug header value. @@ -873,16 +888,18 @@ class MCPJWTSigner(CustomGuardrail): # FR-9: Debug header # ------------------------------------------------------------------ if self.debug_headers: - new_headers["x-litellm-mcp-debug"] = self._build_debug_header(claims, self._kid) + debug_claims: Final[_DebugHeaderClaims] = claims + new_headers["x-litellm-mcp-debug"] = self._build_debug_header(debug_claims, self._kid) hook_data["extra_headers"] = new_headers + logged_claims: Final[_SignedClaimSummary] = claims verbose_proxy_logger.debug( "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d verified=%s channel=%s call_type=%s", - claims.get("sub"), - claims.get("act", {}).get("sub"), + logged_claims.get("sub"), + logged_claims.get("act", {}).get("sub"), hook_data.get("mcp_tool_name"), - claims["exp"], + logged_claims["exp"], jwt_claims is not None, bool(self.channel_token_audience), call_type, diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 704e2564ef5..292f395053b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage +from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus @@ -83,7 +84,8 @@ class NomaV2Guardrail(CustomGuardrail): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - super().__init__(**kwargs) + base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs + super().__init__(**base_kwargs) @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: @@ -114,7 +116,7 @@ class NomaV2Guardrail(CustomGuardrail): return parsed.hostname == _DEFAULT_API_BASE_HOSTNAME @staticmethod - def _get_non_empty_str(value: Any) -> str | None: + def _get_non_empty_str(value: object) -> str | None: if not isinstance(value, str): return None stripped: Final = value.strip() @@ -156,7 +158,7 @@ class NomaV2Guardrail(CustomGuardrail): else model_call_details ) - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "inputs": inputs, "request_data": payload_request_data, "input_type": input_type, @@ -324,8 +326,9 @@ class NomaV2Guardrail(CustomGuardrail): except NomaBlockedMessage as e: guardrail_status = "guardrail_intervened" + blocked_detail: Final[dict[str, object]] = {"error": "blocked"} guardrail_json_response = ( - response_json if isinstance(response_json, dict) else getattr(e, "detail", {"error": "blocked"}) + response_json if isinstance(response_json, dict) else getattr(e, "detail", blocked_detail) ) raise except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 78639ce4fd0..7021d41475b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -8,11 +8,12 @@ # Standard library imports import json import os -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol from urllib.parse import quote # Third-party imports from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict # LiteLLM imports from litellm import DualCache @@ -42,7 +43,34 @@ if TYPE_CHECKING: MAX_PILLAR_HEADER_VALUE_BYTES: Final = 8 * 1024 -def _encode_json_for_header(data: Any) -> str: +class _PillarProtectResponse(TypedDict): + """Body returned by Pillar's `/api/v1/protect` endpoint.""" + + flagged: ReadOnly[NotRequired[bool]] + session_id: ReadOnly[NotRequired[str]] + scanners: ReadOnly[NotRequired[dict[str, object]]] + evidence: ReadOnly[NotRequired[list[object]]] + masked_session_messages: ReadOnly[NotRequired[list[object]]] + + +class _PillarProtectHTTPResponse(Protocol): + def raise_for_status(self) -> object: ... + + def json(self) -> _PillarProtectResponse: ... + + +class _PillarProtectHTTPClient(Protocol): + async def post( + self, + *, + url: str, + headers: dict[str, str], + json: dict[str, object], + timeout: float, + ) -> _PillarProtectHTTPResponse: ... + + +def _encode_json_for_header(data: object) -> str: """ JSON-serialize and URL-encode data for safe header transmission. """ @@ -50,7 +78,9 @@ def _encode_json_for_header(data: Any) -> str: return quote(json_payload, safe="") -def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES) -> tuple[Any, str, bool]: +def _truncate_evidence_payload( + evidence: object, max_bytes: int = MAX_PILLAR_HEADER_VALUE_BYTES +) -> tuple[object, str, bool]: """ Truncate evidence payload so the encoded header value stays within max_bytes. @@ -66,12 +96,12 @@ def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER truncated_value: Final = "[truncated]" return truncated_value, _encode_json_for_header(truncated_value), True - truncated: Final[list[Any]] = [] + truncated: Final[list[object]] = [] encoded = _encode_json_for_header(truncated) truncated_flag = False for entry in evidence: - working_entry: Any + working_entry: object if isinstance(entry, dict): working_entry = dict(entry) else: @@ -105,7 +135,7 @@ def _truncate_evidence_payload(evidence: Any, max_bytes: int = MAX_PILLAR_HEADER return truncated, encoded, truncated_flag -def build_pillar_response_headers(metadata_store: dict[str, Any]) -> dict[str, str]: +def build_pillar_response_headers(metadata_store: dict[str, object]) -> dict[str, str]: """ Create URL-safe Pillar response headers and apply truncation metadata. """ @@ -191,7 +221,9 @@ class PillarGuardrail(CustomGuardrail): LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler: _PillarProtectHTTPClient = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") if self.api_key is None: @@ -686,7 +718,7 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> dict[str, Any]: + async def _call_pillar_api(self, headers: dict[str, str], payload: dict[str, Any]) -> _PillarProtectResponse: """ Call the Pillar API and return the response. @@ -714,7 +746,7 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Analysis complete - flagged=%s, session=%s", flagged, session_id) return res - def _process_pillar_response(self, pillar_response: dict[str, Any], original_data: dict) -> None: + def _process_pillar_response(self, pillar_response: _PillarProtectResponse, original_data: dict) -> None: """ Process the Pillar API response and handle detections based on configuration. @@ -774,7 +806,7 @@ class PillarGuardrail(CustomGuardrail): build_pillar_response_headers(metadata_store) - def _raise_pillar_detection_exception(self, pillar_response: dict[str, Any]) -> None: + def _raise_pillar_detection_exception(self, pillar_response: _PillarProtectResponse) -> None: """ Raise an HTTPException for Pillar security detections. @@ -784,7 +816,7 @@ class PillarGuardrail(CustomGuardrail): Raises: HTTPException: Always raises with security detection details """ - pillar_response_dict: Final = { + pillar_response_dict: Final[dict[str, object]] = { "session_id": pillar_response.get("session_id"), } diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 8d78393d687..da51a905ae3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -11,10 +11,10 @@ import asyncio import json import threading -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Sequence from contextlib import asynccontextmanager from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast import aiohttp from typing_extensions import NotRequired, ReadOnly @@ -68,6 +68,14 @@ class _PresidioAnonymizeResponse(TypedDict): items: ReadOnly[NotRequired[list[_PresidioAnonymizeItem]]] +class _JsonResponse(Protocol): + def json(self) -> Awaitable[object]: ... + + +async def _json_body(response: _JsonResponse) -> object: + return await response.json() + + _LoopSemaphores = dict[asyncio.AbstractEventLoop, asyncio.Semaphore] @@ -389,7 +397,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): f"expected application/json Content-Type but received '{content_type}'; body: '{error_body[:200]}'" ) - analyze_results: Final = await response.json() + analyze_results: Final = await _json_body(response) verbose_proxy_logger.debug("analyze_results: %s", analyze_results) # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) @@ -997,7 +1005,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: raise e - def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: from concurrent.futures import ThreadPoolExecutor def run_in_new_loop(): @@ -1025,7 +1033,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): # No running event loop, we can safely run in this thread return run_in_new_loop() - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: """ Masks the input and output before logging to langfuse, datadog, etc. """ @@ -1092,9 +1100,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): and not isinstance(result.choices[0], StreamingChoices) ): await self._process_response_for_pii(response=result, request_data=kwargs, mode="mask") - elif self._is_anthropic_message_response(result): + elif isinstance(result, dict) and self._is_anthropic_message_response(result): await self._process_anthropic_response_for_pii( - response=cast(dict, result), # cast-ok: _is_anthropic_message_response narrows via isinstance + response=result, request_data=kwargs, mode="mask", ) @@ -1321,7 +1329,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _stream_apply_output_masking( self, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """Apply Presidio masking to streaming output (apply_to_output=True path).""" @@ -1425,7 +1433,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return "\n".join(result_lines).encode("utf-8") - def _unmask_responses_api_completed_chunk(self, chunk: Any, pii_tokens: dict[str, str]) -> None: + def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None: """ Unmask PII tokens in-place for a ``response.completed`` Responses API event. @@ -1434,7 +1442,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): blocks; text blocks expose a ``.text`` string attribute. We walk the tree and replace every PII token with its original value. """ - response_obj: Final = getattr(chunk, "response", None) + response_obj: Final[object] = getattr(chunk, "response", None) if response_obj is None: return @@ -1450,7 +1458,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def _stream_pii_unmasking( self, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """Apply PII unmasking to streaming output (output_parse_pii=True path).""" @@ -1526,7 +1534,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream | bytes, None]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index e34beec4d3e..2fbd50b5863 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -6,7 +6,7 @@ via embedding similarity. Smarter than regex (understands intent), lighter than an LLM call (~20-50ms per request for embedding). """ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from litellm._logging import verbose_logger from litellm.integrations.custom_guardrail import ( @@ -50,7 +50,7 @@ class SemanticGuardrail(CustomGuardrail): similarity_threshold: float, route_templates: list[str] | None = None, custom_routes_file: str | None = None, - custom_routes: list[dict[str, Any]] | None = None, + custom_routes: list[dict[str, object]] | None = None, on_flagged_action: str = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, default_on: bool = False, @@ -157,7 +157,14 @@ class SemanticGuardrail(CustomGuardrail): return response -def _get_top_route_choice(result: Any) -> Any: +class _RouteChoice(Protocol): + """The semantic-router match this guardrail reads: the route that fired, if any.""" + + @property + def name(self) -> str | None: ... + + +def _get_top_route_choice(result: _RouteChoice | list[_RouteChoice] | None) -> _RouteChoice | None: """Extract the top RouteChoice from SemanticRouter result. SemanticRouter.__call__ can return RouteChoice or List[RouteChoice]. @@ -194,7 +201,7 @@ def _extract_response_text(response: Any) -> str: return "" -def _content_to_text(content: Any) -> str: +def _content_to_text(content: object) -> str: if isinstance(content, str): return content if isinstance(content, list): diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 3c5625bc272..a8b33109900 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,9 +1,10 @@ import json import re from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict from fastapi import HTTPException +from typing_extensions import ReadOnly, Required from litellm import ChatCompletionToolParam from litellm._logging import verbose_proxy_logger @@ -51,6 +52,27 @@ def _object_list(value: object) -> Sequence[object] | None: return value if isinstance(value, list) else None +class _ToolPermissionRuleFields(TypedDict, total=False): + """The config-file shape a :class:`ToolPermissionRule` is built from.""" + + id: ReadOnly[Required[str]] + tool_name: ReadOnly[str | None] + tool_type: ReadOnly[str | None] + decision: ReadOnly[Required[Literal["allow", "deny"]]] + allowed_param_patterns: ReadOnly[dict[str, str] | None] + + +def _rule_from_fields(fields: _ToolPermissionRuleFields) -> ToolPermissionRule: + """Validate one config-file rule entry into a :class:`ToolPermissionRule`.""" + return ToolPermissionRule(**fields) + + +def _is_tool_use_block(block: object) -> bool: + """Whether ``block`` is an Anthropic ``tool_use`` content block.""" + fields: Final = _object_mapping(block) + return fields is not None and fields.get("type") == "tool_use" + + class ToolPermissionGuardrail(CustomGuardrail): def __init__( self, @@ -101,7 +123,7 @@ class ToolPermissionGuardrail(CustomGuardrail): compiled_patterns: Final[dict[str, dict[str, re.Pattern]]] = {} for rule_item in rules or []: - rule = rule_item if isinstance(rule_item, ToolPermissionRule) else ToolPermissionRule(**rule_item) + rule = rule_item if isinstance(rule_item, ToolPermissionRule) else _rule_from_fields(rule_item) target_patterns: dict[str, re.Pattern | None] = { "tool_name": None, @@ -440,7 +462,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return is_allowed, None, message @staticmethod - def _get_mapping_value(item: Any, key: str) -> Any: + def _get_mapping_value(item: object, key: str) -> Any: if isinstance(item, dict): return item.get(key) return getattr(item, key, None) @@ -450,7 +472,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return f"legacy_function_call_{choice_index}" def _legacy_function_call_to_tool_call( - self, function_call: Any, choice_index: int + self, function_call: object, choice_index: int ) -> ChatCompletionMessageToolCall | None: if function_call is None: return None @@ -549,7 +571,7 @@ class ToolPermissionGuardrail(CustomGuardrail): def _modify_anthropic_content_with_permission_errors( self, response: object, - content: tuple[Any, ...], + content: tuple[object, ...], denied_tools: tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...], ) -> None: if not denied_tools or not isinstance(response, dict): @@ -557,27 +579,33 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tools)) - error_by_tool_use_id: Final = { # mutable-ok: read-only lookup, never mutated after construction + error_by_tool_use_id: Final[ + Mapping[object, str] + ] = { # mutable-ok: read-only lookup, never mutated after construction tool_call.id: self._create_permission_error_result(tool_call, error).content for tool_call, error in denied_tools } - denied_block_ids: Final = frozenset(error_by_tool_use_id) - def _is_denied(block: object) -> bool: - return isinstance(block, dict) and block.get("type") == "tool_use" and block.get("id") in denied_block_ids + def _denied_message(block: object) -> str | None: + fields: Final = _object_mapping(block) + if fields is None or fields.get("type") != "tool_use": + return None + return error_by_tool_use_id.get(fields.get("id")) - error_messages: Final = tuple(error_by_tool_use_id[block["id"]] for block in content if _is_denied(block)) - kept_blocks: Final = tuple(block for block in content if not _is_denied(block)) + error_messages: Final = tuple( + message for message in (_denied_message(block) for block in content) if message is not None + ) + kept_blocks: Final = tuple(block for block in content if _denied_message(block) is None) new_content: Final = [ # mutable-ok: response content is a JSON array on the wire *kept_blocks, {"type": "text", "text": "\n".join(error_messages)}, # mutable-ok: content block is a JSON object ] response["content"] = new_content # rebind-ok: the guardrail rewrites the provider response in place - if not any(isinstance(block, dict) and block.get("type") == "tool_use" for block in kept_blocks): + if not any(_is_tool_use_block(block) for block in kept_blocks): response["stop_reason"] = "end_turn" # rebind-ok: dropping every tool_use ends the turn - def _get_request_tool_name(self, tool: Any) -> tuple[str | None, str | None]: + def _get_request_tool_name(self, tool: object) -> tuple[str | None, str | None]: tool_type: Final = self._get_mapping_value(tool, "type") if tool_type != "function": return None, tool_type @@ -586,7 +614,7 @@ class ToolPermissionGuardrail(CustomGuardrail): tool_name: Final = self._get_mapping_value(function, "name") return tool_name, tool_type - def _get_legacy_function_name(self, function: Any) -> str | None: + def _get_legacy_function_name(self, function: object) -> str | None: return self._get_mapping_value(function, "name") def _get_named_tool_choice(self, data: dict) -> str | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index 6b8148645aa..a5945a39589 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -433,7 +433,7 @@ class VigilGuardGuardrail(CustomGuardrail): return collected @staticmethod - def _clamp_metadata_value(value: Any) -> _MetadataValue | None: + def _clamp_metadata_value(value: object) -> _MetadataValue | None: if isinstance(value, bool): return None if isinstance(value, str): diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index ddb40dc3ca0..831df43692b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -310,7 +310,7 @@ class XecGuardGuardrail(CustomGuardrail): scan_type: str, suppress_errors: bool = False, ) -> dict | None: - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "model": self.xecguard_model, "scan_type": scan_type, "messages": messages, @@ -385,7 +385,7 @@ class XecGuardGuardrail(CustomGuardrail): def _build_full_history( self, request_data: dict, - inputs: Any, + inputs: GenericGuardrailAPIInputs, input_type: str, ) -> list[dict]: """Assemble the full message list that will be sent to XecGuard. diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 569ec32c1a0..9edbc6dbf1c 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -67,6 +67,24 @@ class _ChatMessage(Protocol): def tool_calls(self) -> Sequence[_ChatToolCall] | None: ... +class _ChatChoice(Protocol): + @property + def message(self) -> _ChatMessage: ... + + @property + def finish_reason(self) -> str | None: ... + + +class _ChatCompletion(Protocol): + @property + def choices(self) -> Sequence[_ChatChoice]: ... + + +def _first_choice(response: _ChatCompletion) -> _ChatChoice: + """The first choice of an OpenAI shaped completion response.""" + return response.choices[0] + + class SkillsInjectionHook(CustomLogger): """ Pre/Post-call hook that processes skills from container.skills parameter. @@ -738,8 +756,9 @@ print('No executable skill module found') for iteration in range(self.max_iterations): # OpenAI format response has choices[0].message - assistant_message: _ChatMessage = current_response.choices[0].message - stop_reason: str | None = current_response.choices[0].finish_reason + choice: _ChatChoice = _first_choice(current_response) + assistant_message: _ChatMessage = choice.message + stop_reason: str | None = choice.finish_reason # Build assistant message for conversation history assistant_msg_dict: dict[str, object] = { diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 3ce406eef73..8b82842353c 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,10 +5,11 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Optional +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -30,6 +31,13 @@ if TYPE_CHECKING: from litellm.router import Router +class SemanticToolFilterConfig(TypedDict, total=False): + enabled: ReadOnly[bool] + embedding_model: ReadOnly[str] + top_k: ReadOnly[int] + similarity_threshold: ReadOnly[float] + + def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> str: """Cap a CSV of tool names to max_length, dropping any name that does not fit whole.""" if len(tool_names_csv) <= max_length: @@ -68,7 +76,7 @@ class SemanticToolFilterHook(CustomLogger): semantic_filter.top_k, ) - def _should_expand_mcp_tools(self, tools: list[Any]) -> bool: + def _should_expand_mcp_tools(self, tools: Iterable[Mapping[str, object]]) -> bool: """ Check if tools contain MCP references with server_url="litellm_proxy". @@ -82,9 +90,9 @@ class SemanticToolFilterHook(CustomLogger): async def _expand_mcp_tools( self, - tools: list[Any], + tools: Iterable[Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth", - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Expand MCP references to actual tool definitions. @@ -111,7 +119,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Convert Pydantic models to dicts for compatibility - openai_tools_as_dicts: Final = [] + openai_tools_as_dicts: Final[list[dict[str, object]]] = [] for tool in openai_tools: if hasattr(tool, "model_dump"): tool_dict = tool.model_dump(exclude_none=True) @@ -141,8 +149,8 @@ class SemanticToolFilterHook(CustomLogger): async def _filter_expanded_tools( self, data: dict, - expanded_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + expanded_tools: list[dict[str, object]], + ) -> list[dict[str, object]]: """ Apply the semantic filter to expanded MCP tool definitions. @@ -159,7 +167,7 @@ class SemanticToolFilterHook(CustomLogger): return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools) - def _selected_tool_names(self, filtered_tools: list[dict[str, Any]]) -> list[str]: + def _selected_tool_names(self, filtered_tools: Sequence[object]) -> list[str]: """Names of the semantically selected tools, as produced by the MCP expansion.""" names: Final = (self.filter._extract_tool_info(tool)[0] for tool in filtered_tools) return [name for name in names if name] @@ -217,10 +225,10 @@ class SemanticToolFilterHook(CustomLogger): def _emit_filter_metadata( self, data: dict, - mcp_tools: list[object], - filtered_mcp_tools: list[object], - native_tools: list[object], - filtered_tools: list[object], + mcp_tools: Sequence[object], + filtered_mcp_tools: Sequence[object], + native_tools: Sequence[object], + filtered_tools: Sequence[object], ) -> None: """ Emit response-header metadata when MCP tools were filtered. @@ -252,10 +260,10 @@ class SemanticToolFilterHook(CustomLogger): def _emit_filter_metadata_safe( self, data: dict, - mcp_tools: list[object], - filtered_mcp_tools: list[object], - native_tools: list[object], - filtered_tools: list[object], + mcp_tools: Sequence[object], + filtered_mcp_tools: Sequence[object], + native_tools: Sequence[object], + filtered_tools: Sequence[object], ) -> None: """ Emit filter metadata without letting an emission failure abort the @@ -375,7 +383,7 @@ class SemanticToolFilterHook(CustomLogger): ) if mcp_tools: - filtered_mcp_tools = await self.filter.filter_tools( + filtered_mcp_tools: list[object] = await self.filter.filter_tools( query=user_query, available_tools=mcp_tools, ) @@ -419,9 +427,9 @@ class SemanticToolFilterHook(CustomLogger): self, data: dict, user_api_key_dict: "UserAPIKeyAuth", - response: Any, + response: object, request_headers: dict[str, str] | None = None, - litellm_call_info: dict[str, Any] | None = None, + litellm_call_info: dict[str, object] | None = None, ) -> dict[str, str] | None: """Add semantic filter stats and tool names to response headers.""" from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH @@ -446,7 +454,7 @@ class SemanticToolFilterHook(CustomLogger): return headers - def _get_tool_names_csv(self, tools: list[Any]) -> str: + def _get_tool_names_csv(self, tools: Sequence[object]) -> str: """Extract tool names and return as CSV string.""" if not tools: return "" @@ -461,7 +469,7 @@ class SemanticToolFilterHook(CustomLogger): @staticmethod async def initialize_from_config( - config: dict[str, Any] | None, + config: SemanticToolFilterConfig | None, llm_router: Optional["Router"], ) -> Optional["SemanticToolFilterHook"]: """ diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 1e65da5b867..63129602082 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -8,7 +8,7 @@ import asyncio import binascii import os import uuid -from collections.abc import Callable, Mapping, Sequence, Set +from collections.abc import Awaitable, Callable, Mapping, Sequence, Set from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime @@ -386,6 +386,12 @@ CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None] ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes +class _AsyncLuaScript(Protocol): + """A Lua script registered against the async Redis client, called with KEYS and ARGV.""" + + def __call__(self, *, keys: Sequence[str], args: Sequence[object]) -> Awaitable[list[CacheCounterValue]]: ... + + class RateLimitDescriptorRateLimitObject(TypedDict, total=False): requests_per_unit: int | None tokens_per_unit: int | None @@ -577,6 +583,14 @@ def _parse_output_cap_value(raw_value: object) -> int | None: class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): + batch_rate_limiter_script: _AsyncLuaScript | None + token_increment_script: _AsyncLuaScript | None + check_and_increment_by_n_script: _AsyncLuaScript | None + window_guarded_token_increment_script: _AsyncLuaScript | None + parallel_acquire_script: _AsyncLuaScript | None + parallel_release_script: _AsyncLuaScript | None + parallel_count_script: _AsyncLuaScript | None + def __init__( self, internal_usage_cache: InternalUsageCache, @@ -3855,7 +3869,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): expected_window_start = operation.get("expected_window_start") if window_key is None or expected_window_start is None: continue - active_window_start = await self.internal_usage_cache.async_get_cache( + active_window_start: CacheCounterValue | None = await self.internal_usage_cache.async_get_cache( key=window_key, litellm_parent_otel_span=parent_otel_span, local_only=True, @@ -4144,7 +4158,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _collect_tpm_scope_targets( self, standard_logging_metadata: dict[str, Any], - kwargs: Any, + kwargs: object, model_group: str | None, ) -> list[tuple[str, str]]: """ @@ -4301,8 +4315,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_success_event_pipeline_operations( self, - kwargs: Any, - response_obj: Any, + kwargs: dict[str, Any], + response_obj: object, rate_limit_type: Literal["output", "input", "total"], ) -> list[RedisPipelineIncrementOperation]: """Build Redis pipeline increment ops for TPM / parallel-request counters.""" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ae55b7ab906..20f83085286 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4,9 +4,10 @@ import json import re import time from collections import OrderedDict -from collections.abc import Mapping, MutableMapping +from collections.abc import Mapping, MutableMapping, Sequence +from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException, Request from pydantic import ValidationError as PydanticValidationError @@ -55,7 +56,7 @@ from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_head from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY # Cache special headers as a frozenset for O(1) lookup performance -_SPECIAL_HEADERS_CACHE: Final = frozenset(v.value.lower() for v in SpecialHeaders._member_map_.values()) +_SPECIAL_HEADERS_CACHE: Final = frozenset(str(v.value).lower() for v in SpecialHeaders) _REDACTED_HEADER_VALUE: Final = "***REDACTED***" _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset( @@ -126,7 +127,7 @@ def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: _ANTHROPIC_SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]+$") -def _sanitize_for_log(value: Any) -> str: +def _sanitize_for_log(value: object) -> str: """ Basic log sanitization helper to reduce log-injection risk. @@ -164,7 +165,7 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig - from litellm.types.proxy.policy_engine import PolicyMatchContext + from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext ProxyConfig = _ProxyConfig else: @@ -328,7 +329,7 @@ _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_overr _URL_DESTINATION_REQUEST_FIELDS: Final = ("model", "file_id") -def _reject_url_valued_destinations(data: dict[str, Any]) -> None: +def _reject_url_valued_destinations(data: dict[str, object]) -> None: """Reject URL-valued ``model``/``file_id`` unless admin-allowlisted. Some providers (HuggingFace, Oobabooga, Gemini files) accept a URL in the @@ -387,7 +388,7 @@ def _invalid_metadata_type_error(field: str, value: object) -> ProxyException: ) -def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]: +def _normalized_metadata_object(field: str, value: object) -> Mapping[str, object]: """Return ``value`` as a metadata object or raise a 400 like OpenAI does. A JSON string that parses to an object is accepted because multipart/form-data @@ -402,6 +403,23 @@ def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]: raise _invalid_metadata_type_error(field=field, value=value) +def _normalized_metadata_slot( + request_data: MutableMapping[str, object], metadata_variable_name: str +) -> dict[str, object]: + """Return the request's metadata slot as a dict, normalising it in place first. + + Metadata can arrive as a JSON string (multipart/form-data, ``extra_body``). Parsing it here keeps + existing entries alive through a merge instead of silently overwriting them with an empty dict. + """ + raw: Final = request_data.get(metadata_variable_name) + if isinstance(raw, dict): + return raw + parsed: Final = safe_json_loads(raw) if isinstance(raw, str) else None + normalized: Final[dict[str, object]] = parsed if isinstance(parsed, dict) else {} + request_data[metadata_variable_name] = normalized + return normalized + + def _strip_untrusted_request_header_controls( headers: Any, *, @@ -417,7 +435,7 @@ def _strip_untrusted_request_header_controls( headers.pop(header_name, None) -def _is_false_like(value: Any) -> bool: +def _is_false_like(value: object) -> bool: if isinstance(value, bool): return value is False if isinstance(value, str): @@ -462,7 +480,7 @@ def _key_or_team_allows_client_pricing_override( ) -def _strip_client_message_redaction_opt_out(data: dict[str, Any]) -> None: +def _strip_client_message_redaction_opt_out(data: dict[str, object]) -> None: stripped: Final[list[str]] = [] if "turn_off_message_logging" in data and _is_false_like(data["turn_off_message_logging"]): stripped.append("turn_off_message_logging") @@ -513,7 +531,7 @@ def _strip_client_callback_credentials( ) -def _strip_client_pricing_overrides(data: dict[str, Any]) -> None: +def _strip_client_pricing_overrides(data: dict[str, object]) -> None: """Drop pricing overrides from the request body and any metadata variant. Skipped only when the calling key/team carries @@ -580,9 +598,9 @@ def _get_metadata_variable_name(request: Request) -> str: def _promoted_trace_control_fields( - requester_metadata: Mapping[str, Any], - litellm_metadata: Mapping[str, Any], -) -> tuple[tuple[str, Any], ...]: + requester_metadata: Mapping[str, object], + litellm_metadata: Mapping[str, object], +) -> tuple[tuple[str, object], ...]: """Return the caller's trace-control fields that ``litellm_metadata`` does not already set.""" return tuple( (key, value) @@ -1193,7 +1211,7 @@ class LiteLLMProxyRequestSetup: def add_litellm_data_for_backend_llm_call( *, headers: dict, - request_data: Mapping[str, Any], + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, general_settings: dict[str, Any] | None = None, ) -> LitellmDataForBackendLLMCall: @@ -1327,6 +1345,8 @@ class LiteLLMProxyRequestSetup: def get_sanitized_user_information_from_key( user_api_key_dict: UserAPIKeyAuth, ) -> StandardLoggingUserAPIKeyMetadata: + stripped_metadata: Final = strip_callback_config(user_api_key_dict.metadata) + auth_metadata: Final = cast("dict[str, str] | None", stripped_metadata) # cast-ok: metadata is free-form JSON user_api_key_logged_metadata: Final = StandardLoggingUserAPIKeyMetadata( user_api_key_hash=user_api_key_dict.api_key, # just the hashed token user_api_key_alias=user_api_key_dict.key_alias, @@ -1349,7 +1369,7 @@ class LiteLLMProxyRequestSetup: user_api_key_budget_reset_at=( user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), - user_api_key_auth_metadata=strip_callback_config(user_api_key_dict.metadata), + user_api_key_auth_metadata=auth_metadata, ) return user_api_key_logged_metadata @@ -1577,14 +1597,7 @@ class LiteLLMProxyRequestSetup: return _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata = request_data.get(_metadata_variable_name) - if isinstance(metadata, str): - parsed: Final = safe_json_loads(metadata) - metadata = parsed if isinstance(parsed, dict) else {} - request_data[_metadata_variable_name] = metadata - elif not isinstance(metadata, dict): - metadata = {} - request_data[_metadata_variable_name] = metadata + metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( @@ -1636,18 +1649,7 @@ class LiteLLMProxyRequestSetup: # from (litellm_metadata vs metadata) so the merged tags are visible # to _tag_max_budget_check. _metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data) - metadata = request_data.get(_metadata_variable_name) - # metadata can arrive as a JSON string (multipart/form-data, extra_body). - # Parse it so existing tags survive the merge — overwriting the string - # with {} would let a caller bypass _tag_max_budget_check on an - # over-budget body tag by also sending a within-budget header tag. - if isinstance(metadata, str): - parsed: Final = safe_json_loads(metadata) - metadata = parsed if isinstance(parsed, dict) else {} - request_data[_metadata_variable_name] = metadata - elif not isinstance(metadata, dict): - metadata = {} - request_data[_metadata_variable_name] = metadata + metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name) existing_tags: Final = metadata.get("tags") metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( @@ -1787,7 +1789,7 @@ async def add_litellm_data_to_request( # admin-injection strip below so the audit / spend-tracking consumers of # proxy_server_request["body"] see the cleaned metadata rather than # attacker-forged user_api_key_* fields. - _litellm_received_at: Final = getattr(request.state, "litellm_received_at", None) + _litellm_received_at: Final[datetime | None] = getattr(request.state, "litellm_received_at", None) arrival_time: Final = _litellm_received_at.timestamp() if _litellm_received_at is not None else time.time() data["proxy_server_request"] = { "url": str(request.url), @@ -2472,16 +2474,16 @@ def _resolve_provider_from_deployment( if deployment is None: continue - litellm_params = getattr(deployment, "litellm_params", None) + litellm_params: object = getattr(deployment, "litellm_params", None) if litellm_params is None: continue custom_provider = getattr(litellm_params, "custom_llm_provider", None) - if custom_provider: + if isinstance(custom_provider, str) and custom_provider: return custom_provider - deployment_model = getattr(litellm_params, "model", "") or "" - if "/" in deployment_model: + deployment_model = getattr(litellm_params, "model", "") + if isinstance(deployment_model, str) and "/" in deployment_model: return deployment_model.split("/", 1)[0] return None @@ -2904,8 +2906,8 @@ def _extract_policy_id(s: str) -> str | None: def _match_and_track_policies( data: dict, context: "PolicyMatchContext", - request_body_policies: Any, - policies_override: dict[str, Any] | None = None, + request_body_policies: Sequence[str], + policies_override: dict[str, "Policy"] | None = None, ) -> tuple[list[str], dict[str, str]]: """ Match policies via attachments and request body, track them in metadata. @@ -2963,7 +2965,7 @@ def _apply_resolved_guardrails_to_metadata( metadata_variable_name: str, context: "PolicyMatchContext", policy_names: list[str] | None = None, - policies: dict[str, Any] | None = None, + policies: dict[str, "Policy"] | None = None, ) -> None: """Apply resolved guardrails and pipelines to request metadata.""" from litellm._logging import verbose_proxy_logger @@ -3093,7 +3095,7 @@ async def add_guardrails_from_policy_engine( request_body_names.append(item) # Resolve policy versions by ID from in-memory cache (populated by sync job; no DB in hot path) - merged_policies: Final[dict[str, Any]] = dict(registry.get_all_policies()) + merged_policies: Final[dict[str, Policy]] = dict(registry.get_all_policies()) fetched_policy_names: Final[list[str]] = [] for policy_id in request_body_version_ids: result = registry.get_policy_by_id_for_request(policy_id=policy_id) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 44b0cdcca2e..21e652114bc 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -833,7 +833,7 @@ def _judge_collisions_for_team( return tuple( (role, model) for role, model in ( - *_router_arm_models(llm_router, data.router_name), + *(arm for name in data.router_names for arm in _router_arm_models(llm_router, name)), *((("baseline", data.baseline_model),) if data.baseline_model is not None else ()), ) if judge & judge_target(llm_router, model, team_id).models @@ -904,7 +904,7 @@ class _AttemptAggRow(BaseModel): _ATTEMPT_AGG_ROWS: Final = TypeAdapter(list[_AttemptAggRow]) -_ATTEMPT_AGG_SELECT: Final = """ +_ATTEMPT_AGG_COLUMNS: Final = """ COUNT(*)::int AS turn_count, COUNT(*) FILTER (WHERE outcome = 'real')::int AS real_wins, COUNT(*) FILTER (WHERE outcome = 'shadow')::int AS shadow_wins, @@ -913,15 +913,34 @@ _ATTEMPT_AGG_SELECT: Final = """ COALESCE(SUM(real_cost + real_classifier_cost) FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit), 0)::float AS real_spend, COALESCE(SUM(shadow_cost + shadow_classifier_cost) FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit), 0)::float AS shadow_spend, COUNT(*) FILTER (WHERE real_cache_hit)::int AS cache_hit_turns +""" + +_ATTEMPT_AGG_SELECT: Final = ( + _ATTEMPT_AGG_COLUMNS + + """ FROM "LiteLLM_ShadowEvalAttempt" WHERE job_id = ANY($1::text[]) AND outcome != 'error' GROUP BY 1 """ +) _ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT +# Attempt rows from before arm stamping carry no router_name; they belong to the job's +# own router, which the join reads off the leg. +_ATTEMPT_AGG_BY_ROUTER_SQL: Final = ( + "SELECT COALESCE(a.router_name, j.router_name) AS grp," + + _ATTEMPT_AGG_COLUMNS + + """ +FROM "LiteLLM_ShadowEvalAttempt" a +JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id +WHERE a.job_id = ANY($1::text[]) AND a.outcome != 'error' +GROUP BY 1 +""" +) + # These guards derive spend from attempt rows, the cross-pod authority; the sampler also # reads the live counter, so admission can stop before a row-based guard would fire (safe # direction, and mid-deploy rows from old pods price as judge-only until the deploy ends). @@ -1060,6 +1079,7 @@ class _LegRow(BaseModel): target_type: ShadowEvalTargetType target_id: str router_name: str + router_names: tuple[str, ...] = () direction: ShadowEvalDirection baseline_model: str | None = None judge_model: str @@ -1071,6 +1091,12 @@ class _LegRow(BaseModel): stopped_at: datetime | None = None stopped_by: str | None = None + @property + def arm_router_names(self) -> tuple[str, ...]: + """The job's full router set; rows from before router_names existed hold it in + router_name alone. The one place that reading lives on the endpoint side.""" + return self.router_names or (self.router_name,) + @field_validator("created_at", "ends_at", "stopped_at") @classmethod def _as_aware_utc(cls, value: datetime | None) -> datetime | None: @@ -1123,7 +1149,7 @@ def _group_response( ) for leg in sorted(legs, key=lambda leg: (leg.target_type, leg.target_id)) ), - router_name=first.router_name, + router_names=first.arm_router_names, direction=first.direction, baseline_model=first.baseline_model, judge_model=first.judge_model, @@ -1252,6 +1278,9 @@ async def _shadow_eval_results( for slice in _slices(by_leg) } ) + by_router: Final = _ATTEMPT_AGG_ROWS.validate_python( + await _query_raw(prisma_client, _ATTEMPT_AGG_BY_ROUTER_SQL, leg_ids) or () + ) total_turns: Final = sum(r.turn_count for r in by_tier) funnel_rows: Final = await _query_raw(prisma_client, _FUNNEL_TOTALS_SQL, leg_ids) counted: Final = _FunnelTotalsRow.model_validate(funnel_rows[0]) if funnel_rows else None @@ -1261,6 +1290,7 @@ async def _shadow_eval_results( result: Final = ShadowEvalResult( by_tier=_slices(by_tier), by_current_model=_slices(by_model), + by_router=_slices(by_router), overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns), overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns), sampled_real_spend=sum(r.real_spend for r in by_tier), @@ -1314,8 +1344,15 @@ async def start_shadow_eval( _require_admin_writer(user_api_key_dict, "start a shadow eval") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, data.router_name): - raise HTTPException(status_code=400, detail=f"'{data.router_name}' is not a configured auto-router") + unconfigured: Final = tuple( + name + for name in data.router_names + if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, name) + ) + if unconfigured: + raise HTTPException( + status_code=400, detail=f"Not a configured auto-router: {', '.join(repr(n) for n in unconfigured)}" + ) token_rows: Final = ( await _verification_tokens(prisma_client).find_many( where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter @@ -1416,7 +1453,9 @@ async def start_shadow_eval( ends_at: Final = now + timedelta(days=data.duration_days) shared_config: Final = { # mutable-ok: Prisma payload "group_id": group_id, - "router_name": data.router_name, + # a pre-router_names pod samples router_name alone, so it must be a real arm + "router_name": data.router_names[0], + "router_names": list(data.router_names), # mutable-ok: Prisma payload "direction": data.direction, "baseline_model": data.baseline_model, "judge_model": data.judge_model, @@ -1477,7 +1516,7 @@ async def start_shadow_eval( ) for target_type, target_id in sorted(requested_targets) ), - router_name=data.router_name, + router_names=data.router_names, direction=data.direction, baseline_model=data.baseline_model, judge_model=data.judge_model, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 26b65dd48cf..c3403cf477c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -20,6 +20,7 @@ import secrets import traceback from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast import fastapi @@ -111,6 +112,7 @@ from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( get_ui_settings_cached, @@ -736,6 +738,45 @@ def _check_allowed_routes_caller_permission( ) +_READ_ONLY_ALLOWED_ROUTES_PRESET: Final = frozenset(("info_routes",)) + + +def _is_safe_preset_route_transition( + incoming_allowed_routes: Sequence[str] | None, + existing_allowed_routes: Sequence[str] | None, +) -> bool: + """ + True when every route on BOTH sides is a safe `key_type` preset bucket + (empty = full access, which non-admins already get from a default + `/key/generate`), with one carve-out: a read-only (`info_routes`) key + stays read-only, so widening it needs an admin. Requiring the existing + side to be a safe preset keeps an owner from clearing an admin-set + custom route restriction (LIT-4139). + """ + incoming: Final = frozenset(incoming_allowed_routes or ()) + existing: Final = frozenset(existing_allowed_routes or ()) + if not (incoming | existing) <= _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS: + return False + return existing != _READ_ONLY_ALLOWED_ROUTES_PRESET or incoming == existing + + +def _enforce_allowed_routes_update_permission( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + if _is_safe_preset_route_transition( + incoming_allowed_routes=data.allowed_routes, + existing_allowed_routes=existing_key_row.allowed_routes, + ): + return + _check_allowed_routes_caller_permission( + allowed_routes=data.allowed_routes, + user_api_key_dict=user_api_key_dict, + allowed_routes_was_provided="allowed_routes" in data.model_fields_set, + ) + + def _check_permissions_caller_permission( data: GenerateRequestBase, user_api_key_dict: UserAPIKeyAuth, @@ -2520,26 +2561,34 @@ async def _validate_mcp_servers_for_key_update( return normalized_object_permission +def _require_prisma_client(prisma_client: PrismaClient | None) -> PrismaClient: + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "Database not connected"}) + return prisma_client + + async def _validate_update_key_data( data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, llm_router: Router | None, premium_user: bool, - prisma_client: Any, + prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ) -> None: """Validate permissions and constraints for key update.""" + checked_prisma_client: Final = _require_prisma_client(prisma_client) + # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) validate_budget_duration(data.budget_duration) _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value - _check_allowed_routes_caller_permission( - allowed_routes=data.allowed_routes, + _enforce_allowed_routes_update_permission( + data=data, + existing_key_row=existing_key_row, user_api_key_dict=user_api_key_dict, - allowed_routes_was_provided="allowed_routes" in data.model_fields_set, ) _check_passthrough_routes_caller_permission( data=data, @@ -2567,7 +2616,7 @@ async def _validate_update_key_data( await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( user_api_key_dict=user_api_key_dict, route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, existing_key_row=existing_key_row, user_api_key_cache=user_api_key_cache, ) @@ -2648,12 +2697,12 @@ async def _validate_update_key_data( # _check_key_admin_access that would otherwise require team/org admin status. _key_is_team_key: Final = getattr(existing_key_row, "team_id", None) is not None can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not _is_budget_change - if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check: + if (not _is_proxy_admin) and not can_skip_admin_check: hashed_key: Final = existing_key_row.token await _check_key_admin_access( user_api_key_dict=user_api_key_dict, hashed_token=hashed_key, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) @@ -2664,7 +2713,7 @@ async def _validate_update_key_data( if _team_id_to_check is not None: team_obj = await get_team_object( team_id=_team_id_to_check, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, ) @@ -2680,7 +2729,7 @@ async def _validate_update_key_data( await _check_team_key_limits( team_table=team_obj, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( @@ -2695,7 +2744,7 @@ async def _validate_update_key_data( await _check_project_key_limits( project_id=_project_id_to_check, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, ) @@ -2710,7 +2759,7 @@ async def _validate_update_key_data( await _validate_caller_can_assign_key_org( user_api_key_dict=user_api_key_dict, organization_id=data.organization_id, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) # Check org key limits only when throughput-related fields or organization_id change @@ -2726,7 +2775,7 @@ async def _validate_update_key_data( org_table: Final = await get_org_object( org_id=_org_id_to_check, user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) if org_table is None: raise HTTPException( @@ -2736,7 +2785,7 @@ async def _validate_update_key_data( await _check_org_key_limits( org_table=org_table, data=data, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, ) # if team change - check if this is possible @@ -2766,7 +2815,7 @@ async def _validate_update_key_data( data=data, team_obj=team_obj, existing_key_row=existing_key_row, - prisma_client=prisma_client, + prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, is_proxy_admin=_is_proxy_admin, ) @@ -3578,6 +3627,63 @@ async def _build_model_max_budget_usage( ) +def _window_max_budget(window: Mapping[str, object]) -> float | None: + """A window's max_budget as a float; None when absent or unparseable.""" + value: Final = window.get("max_budget") + if not isinstance(value, (int, float, str)): + return None + try: + return float(value) + except ValueError: + return None + + +async def _budget_window_usage( + window: Mapping[str, object], api_key_hash: str +) -> tuple[str, Mapping[str, object]] | None: + """ + (budget_duration, usage entry) for one budget window; None when the window + has no budget_duration to key it by. + + Reads the same cross-pod counter (spend:key:{hashed_token}:window:{budget_duration}) + that _virtual_key_multi_budget_check enforces against, passing the same + window_duration + window_start so a stale-low counter is re-checked against + the LiteLLM_BudgetWindowSpend row instead of a spend-log aggregate. + """ + from litellm.proxy.proxy_server import get_current_spend + + duration: Final = window.get("budget_duration") + if not isinstance(duration, str) or not duration: + return None + spend: Final = await get_current_spend( + counter_key=f"spend:key:{api_key_hash}:window:{duration}", + fallback_spend=0.0, + max_budget=_window_max_budget(window), + window_entity_type="Key", + window_entity_id=api_key_hash, + window_duration=duration, + window_start=get_budget_window_start(window), + ) + return duration, MappingProxyType({"current_spend": round(spend, 4)}) + + +async def _build_budget_limits_usage( + budget_limits: Sequence[object] | str | None, api_key_hash: str +) -> Mapping[str, Mapping[str, object]] | None: + """ + Current-window spend per budget window, keyed by budget_duration, reported + next to the stored budget_limits (which is returned untouched). None when + the key has no windows, so the field only appears on keys that have them. + """ + windows: Final = _budget_limit_windows(budget_limits) + if not windows: + return None + usages: Final = await asyncio.gather( + *(_budget_window_usage(window=window, api_key_hash=api_key_hash) for window in windows) + ) + return MappingProxyType({duration: usage for duration, usage in (u for u in usages if u is not None)}) + + @router.post( "/v2/key/info", tags=["key management"], @@ -3620,7 +3726,6 @@ async def info_key_fn_v2( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail={"message": "Malformed request. No keys passed in."}, ) - # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query: Final = list(data.keys) if data.keys else [] if data.key_aliases: @@ -3662,6 +3767,13 @@ async def info_key_fn_v2( model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) + if k_token_hash: + budget_limits_usage = await _build_budget_limits_usage( + budget_limits=k_dict.get("budget_limits"), + api_key_hash=k_token_hash, + ) + if budget_limits_usage is not None: + k_dict["budget_limits_usage"] = budget_limits_usage filtered_key_info.append(k_dict) return {"key": data.keys, "info": filtered_key_info} @@ -3698,6 +3810,10 @@ async def info_key_fn( - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when the key has per-model budgets + - budget_limits: list | None - Concurrent budget windows, exactly as stored + - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. + {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows + (read from the same cross-pod spend counter the budget enforcement uses) - models: list - Model_name's the key is allowed to call - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} @@ -3765,7 +3881,7 @@ async def info_key_fn( except Exception: # if using pydantic v1 key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback - key_token_hash: Final = key_info.pop("token") + key_token_hash: Final[str | None] = key_info.pop("token") model_max_budget = key_info.get("model_max_budget") or {} budget_table: Final = key_info.get("litellm_budget_table") or {} @@ -3777,6 +3893,12 @@ async def info_key_fn( model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) + budget_limits_usage: Final = await _build_budget_limits_usage( + budget_limits=key_info.get("budget_limits"), + api_key_hash=key_token_hash, + ) + if budget_limits_usage is not None: + key_info["budget_limits_usage"] = budget_limits_usage # Attach object_permission if object_permission_id is set key_info = await attach_object_permission_to_dict(key_info, prisma_client) @@ -5221,7 +5343,7 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio max_budget = key_in_db.max_budget if key_in_db.litellm_budget_table is not None: - budget_max_budget: Final = getattr(key_in_db.litellm_budget_table, "max_budget", None) + budget_max_budget: Final[float | None] = getattr(key_in_db.litellm_budget_table, "max_budget", None) if budget_max_budget is not None: if max_budget is None or budget_max_budget < max_budget: max_budget = budget_max_budget diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 012aec38458..ca66640bf46 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -54,7 +54,10 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( @@ -543,7 +546,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) - merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) + merged_model_info: Final[dict[str, object]] = db_model.model_info.model_dump(exclude_none=True) # update litellm params if updated_patch.litellm_params: @@ -701,6 +704,12 @@ async def patch_model( param="blocked", ) + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=patch_data.litellm_params, + user_api_key_dict=user_api_key_dict, + existing_litellm_params=db_model.litellm_params, + ) + _raise_on_strategy_router_write_violation( incoming_params=patch_data.litellm_params, existing_params=db_model.litellm_params, @@ -1464,6 +1473,32 @@ class ModelManagementAuthChecks: ) return True + @staticmethod + def can_user_attach_credential( + litellm_params: GenericLiteLLMParams | None, + user_api_key_dict: UserAPIKeyAuth, + existing_litellm_params: GenericLiteLLMParams | None = None, + ) -> Literal[True]: + if litellm_params is None or litellm_params.litellm_credential_name is None: + return True + if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None: + existing_credential_name: Final = decrypt_value_helper( + value=existing_litellm_params.litellm_credential_name, + key="litellm_credential_name", + exception_type="debug", + return_original_value=True, + ) + if litellm_params.litellm_credential_name == existing_credential_name: + return True + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + raise ProxyException( + message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.", + type=ProxyErrorTypes.auth_error.value, + code=status.HTTP_403_FORBIDDEN, + param="litellm_credential_name", + ) + @staticmethod async def allow_team_model_action( model_params: Deployment | updateDeployment, @@ -1786,6 +1821,11 @@ async def add_new_model( premium_user=premium_user, ) + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=model_params.litellm_params, + user_api_key_dict=user_api_key_dict, + ) + _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=None, @@ -1958,6 +1998,12 @@ async def update_model( premium_user=premium_user, ) + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=model_params.litellm_params, + user_api_key_dict=user_api_key_dict, + existing_litellm_params=deployment.litellm_params, + ) + _raise_on_strategy_router_write_violation( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, @@ -1982,7 +2028,7 @@ async def update_model( ### MERGE WITH EXISTING DATA ### merged_dictionary: Final = {} - _mp: Final = model_params.litellm_params.dict() + _mp: Final[dict[str, object]] = model_params.litellm_params.dict() for key, value in _mp.items(): if value is not None: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 9198aa35f3f..5e38a016099 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -487,12 +487,11 @@ async def new_organization( for m in data.models: await can_user_call_model(m, llm_router=llm_router, user_object=user_object_correct_type) - organization_row: Final = LiteLLM_OrganizationTable( - **data.json(exclude_none=True), - object_permission_id=object_permission_id, - created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, - updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, - ) + organization_payload: Final = _STR_OBJECT_DICT_ADAPTER.validate_python(data.json(exclude_none=True)) + organization_payload["object_permission_id"] = object_permission_id + organization_payload["created_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + organization_payload["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + organization_row: Final = LiteLLM_OrganizationTable.model_validate(organization_payload) for field in LiteLLM_ManagementEndpoint_MetadataFields: if getattr(data, field, None) is not None: @@ -644,7 +643,7 @@ async def update_organization( ) # Transform UI payload to expected format - raw_data: Final = await request.json() + raw_data: Final[dict[str, object]] = await request.json() raw_data_with_flat_budget_fields: Final = handle_nested_budget_structure_in_organization_update_request(raw_data) # Create validated data model @@ -691,7 +690,7 @@ async def update_organization( # Merge metadata from existing organization with updated metadata if updated_organization_row_json.get("metadata") is not None: existing_metadata: Final = existing_organization_row.metadata or {} - updated_metadata: Final = updated_organization_row_json.get("metadata", {}) + updated_metadata: Final[dict[str, object]] = updated_organization_row_json.get("metadata", {}) merged_metadata: Final[Mapping[str, object]] = _update_dictionary( existing_dict=cast( # cast-ok: prisma de-serializes a Json column to the plain python dict it stores "dict[str, object]", existing_metadata diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 108e6a7b47d..f58f3722741 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -13,7 +13,7 @@ import copy import json import os from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import Response, StreamingResponse @@ -90,7 +90,7 @@ class _ApplyPoliciesResultBase(TypedDict): class ApplyPoliciesResult(_ApplyPoliciesResultBase, total=False): """Result of apply_policies. agent_response set when agent_id provided.""" - agent_response: Any + agent_response: object class _ApplyPoliciesPerItemResultBase(TypedDict): @@ -103,7 +103,7 @@ class _ApplyPoliciesPerItemResultBase(TypedDict): class ApplyPoliciesPerItemResult(_ApplyPoliciesPerItemResultBase, total=False): """Result for one input when using inputs_list. agent_response set when agent_id provided.""" - agent_response: Any + agent_response: object class ApplyPoliciesListResult(TypedDict): @@ -295,8 +295,8 @@ async def test_policies_and_guardrails( from litellm.proxy.proxy_server import chat_completion, proxy_logging_obj from litellm.proxy.utils import handle_exception_on_proxy - def _serialize_chat_response(response: Any) -> Any: - if hasattr(response, "model_dump"): + def _serialize_chat_response(response: object) -> object: + if isinstance(response, BaseModel): return response.model_dump(exclude_unset=True) if isinstance(response, dict): return response @@ -306,7 +306,7 @@ async def test_policies_and_guardrails( inputs: GenericGuardrailAPIInputs, agent_id: str, user_api_key_dict: UserAPIKeyAuth, - ) -> Any: + ) -> object: body: Final = _chat_body_from_inputs(inputs, agent_id, data.request_data) req: Final = _request_with_json_body(body) resp: Final = Response() diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index ded57815e91..069f86c852c 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -29,6 +29,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.user import SCIMPlaceholder from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_UserTable, @@ -585,6 +586,37 @@ async def _users_named_by_member_value( return tuple(dict.fromkeys(row.user_id for row in rows)) +async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]: + """Every user id this member value names, by user id, SSO identity or email. + + Classification needs to know whether the value is one account's ``user_id`` and + whether it names any other account, so all three fields are read in one pass. The + id is compared exactly and unstripped, as a primary key lookup would; the + identities compare as ``_users_named_by_member_value`` describes. Two rows are + enough to tell one account from several, so the read stops there. Only a full + read that lacks the row keyed by the value leaves that row's existence open, and + only then is the id read on its own. + """ + subject: Final = value.strip() + email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"} + users: Final = _table(UserRepository(prisma_client)) + rows: Final = await users.find_many( + where={ # mutable-ok: Prisma filter + "OR": [ # mutable-ok: Prisma filter + {"user_id": value}, # mutable-ok: Prisma filter + {"sso_user_id": subject}, # mutable-ok: Prisma filter + {"user_email": email}, # mutable-ok: Prisma filter + ], + }, + take=2, + ) + named: Final = tuple(dict.fromkeys(row.user_id for row in rows)) + if len(named) < 2 or value in named: + return named + keyed: Final = await users.find_unique(where={"user_id": value}) + return named if keyed is None else (value, *named) + + async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember: """ Decide what a single SCIM group member refers to. @@ -627,11 +659,9 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if member_type == "group": return _SkippedGroupMember(value=value, reason="nested_group") - user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value}) - if user is not None: - shared_with: Final = tuple( - other for other in await _users_named_by_member_value(value, prisma_client) if other != value - ) + named: Final = await _accounts_named_by_member_value(value, prisma_client) + if value in named: + shared_with: Final = tuple(other for other in named if other != value) if shared_with: verbose_proxy_logger.warning( "SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, " @@ -651,7 +681,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if team is not None and _team_metadata_has_scim_provenance(team.metadata): return _SkippedGroupMember(value=value, reason="existing_team") - named: Final = await _users_named_by_member_value(value, prisma_client) if len(named) == 1: verbose_proxy_logger.info( "SCIM: group member '%s' matched user_id '%s' by SSO identity or email", @@ -1834,6 +1863,89 @@ async def delete_user( raise handle_exception_on_proxy(e) +@scim_router.get( + "/placeholders", + response_model=tuple[SCIMPlaceholder, ...], + dependencies=(Depends(user_api_key_auth),), +) +async def list_placeholders() -> tuple[SCIMPlaceholder, ...]: + """ + List user rows whose id is another account's SSO identity or email. + + An earlier release provisioned a group member it could not match as a user keyed + by the raw member value, and that row now shadows the account the value really + names, so every push of that member is refused. This lists those rows so an + operator can fold each one into the account it shadows with + ``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of + its own or owns virtual keys is left out: someone uses that account. + """ + try: + prisma_client: Final = await _get_prisma_client_or_raise_exception() + async with prisma_client.tx() as tx: + return await UserRepository(prisma_client).find_shadowing_placeholders(tx) + except Exception as e: + raise handle_exception_on_proxy(e) + + +def _placeholder_rejection(placeholder: LiteLLM_UserTable, resolved: tuple[str, ...], key_count: int) -> str | None: + if placeholder.sso_user_id is not None: + return f"User '{placeholder.user_id}' has an SSO identity of its own, so it is an account someone signs in to" + if key_count: + return f"User '{placeholder.user_id}' owns {key_count} virtual keys. Move or delete them before merging it" + if not resolved: + return f"User '{placeholder.user_id}' shadows no account: no other user has that id as SSO identity or email" + if len(resolved) > 1: + return ( + f"User '{placeholder.user_id}' names {len(resolved)} accounts ({', '.join(resolved)}). Resolve that first" + ) + return None + + +@scim_router.post( + "/placeholders/{user_id}/merge", + response_model=SCIMPlaceholderMergeResult, + dependencies=(Depends(user_api_key_auth),), +) +async def merge_placeholder( + user_id: str = Path(..., title="User ID"), +) -> SCIMPlaceholderMergeResult: + """ + Fold a placeholder user into the one account its id names by SSO identity or email. + + The account is added to every team the placeholder is on, then the placeholder is + deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group + push resolves the member value to the real account. Refused with 409 when the row + has an SSO identity of its own, owns virtual keys, or names no account or several. + """ + try: + prisma_client: Final = await _get_prisma_client_or_raise_exception() + placeholder: Final = await _check_user_exists(user_id) + resolved: Final = tuple( + other for other in await _users_named_by_member_value(user_id, prisma_client, take=None) if other != user_id + ) + owned_keys: Final[_UserIdWhere] = {"user_id": user_id} + keys: Final = await _table(VerificationTokenRepository(prisma_client)).find_many(where=owned_keys) + rejection: Final = _placeholder_rejection(placeholder, resolved, len(keys)) + if rejection is not None: + detail: Final[_ScimErrorDetail] = {"error": rejection} + raise HTTPException(status_code=409, detail=detail) + + target_user_id: Final = resolved[0] + team_ids: Final = tuple(placeholder.teams) + for team_id in team_ids: + await _add_user_to_team(user_id=target_user_id, team_id=team_id) + await delete_user(user_id=user_id) + await _recompute_scim_member_roles(prisma_client, (target_user_id,)) + verbose_proxy_logger.info( + "SCIM: merged placeholder user '%s' into '%s', moving teams %s", user_id, target_user_id, team_ids + ) + return SCIMPlaceholderMergeResult( + placeholder_user_id=user_id, merged_into_user_id=target_user_id, team_ids=team_ids + ) + except Exception as e: + raise handle_exception_on_proxy(e) + + def _parse_member_entry(entry: object) -> SCIMMember | None: """Parse one entry of a SCIM patch value, or None when it carries no id.""" if isinstance(entry, str): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c6d7975b75e..714cf252e69 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -5148,6 +5148,7 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: + # LiteLLM_DeletedTeamTable has no litellm_model_table relation, unlike below teams = await _deleted_team_db(prisma_client).find_many( where=where_conditions, skip=skip, @@ -5162,6 +5163,7 @@ async def list_team_v2( skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort + include=_INCLUDE_MODEL_TABLE, ) # Get total count for pagination total_count = await _team_db(prisma_client).count(where=where_conditions) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 613508da22b..606569c5b8b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -502,7 +502,7 @@ def _set_nested_metadata_value(metadata: dict[str, object], key_path: str, value placeholder: Final = "\x00" parts = key_path.replace("\\.", placeholder).split(".") parts = [p.replace(placeholder, ".") for p in parts] - current: Any = metadata + current: dict[str, object] = metadata for part in parts[:-1]: existing = current.get(part) if not isinstance(existing, dict): @@ -4076,7 +4076,7 @@ class SSOAuthenticationHandler: ) if resp.status_code == 200: try: - userinfo_raw: Final = resp.json() + userinfo_raw: Final[dict[str, object] | None] = resp.json() if not userinfo_raw: # JSON null (None) or empty dict ({}) — no identity claims. # Treat as failure so id_token fallback can be attempted. @@ -4406,7 +4406,7 @@ class MicrosoftSSOHandler: ) -> tuple[list[str], str | None]: """Helper function to fetch and parse group data from a URL""" response: Final = await async_client.get(url, headers=headers) - response_json: Final = response.json() + response_json: Final[dict[str, object]] = response.json() response_typed: Final = await MicrosoftSSOHandler._cast_graph_api_response_dict(response=response_json) group_ids: Final = MicrosoftSSOHandler._get_group_ids_from_graph_api_response(response=response_typed) return group_ids, response_typed.get("odata_nextLink") diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py index 5d43cb29978..c9f93fae0d9 100644 --- a/litellm/proxy/management_helpers/access_group_key_sync.py +++ b/litellm/proxy/management_helpers/access_group_key_sync.py @@ -38,6 +38,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( _delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive ) +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.table_repositories import AccessGroupRepository @@ -72,8 +73,9 @@ _REPOINT_KEY_SQL: Final = ( def _raw_executor(prisma_client: object) -> _RawExecutor: - """Narrow the untyped Prisma client down to the raw-query call this module makes.""" - return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + """Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer.""" + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin async def _invalidate_access_group_cache(access_group_id: str) -> None: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 13080a6cf83..a2fbf80422c 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -286,7 +286,7 @@ async def _resolve_mcp_server_identifiers_to_ids( return resolved -def _rewrite_object_permission_mcp_servers( +def _drop_stale_object_permission_mcp_servers( object_permission: ObjectPermissionDict, identifier_to_server_ids: dict[str, set[str]], ) -> None: @@ -294,16 +294,18 @@ def _rewrite_object_permission_mcp_servers( if not isinstance(mcp_servers, list): return - normalized_servers: Final[list[str]] = [] - for identifier in mcp_servers: - if identifier == SpecialMCPServerNames.no_mcp_servers.value: - normalized_servers.append(SpecialMCPServerNames.no_mcp_servers.value) - continue - normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, []))) - object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers) + # Persist original identifiers, never resolved ids: shared-DB multi-region + # instances each expand a name/alias to their own local server id at read + # time. Only entries resolving to nothing (deleted servers, typos) drop. + kept_servers: Final = [ + identifier + for identifier in mcp_servers + if identifier == SpecialMCPServerNames.no_mcp_servers.value or identifier_to_server_ids.get(identifier) + ] + object_permission["mcp_servers"] = _dedupe_preserving_order(kept_servers) -def _rewrite_object_permission_mcp_tool_permissions( +def _drop_stale_object_permission_mcp_tool_permissions( object_permission: ObjectPermissionDict, identifier_to_server_ids: dict[str, set[str]], ) -> None: @@ -311,31 +313,25 @@ def _rewrite_object_permission_mcp_tool_permissions( if not isinstance(mcp_tool_permissions, dict): return - normalized_tool_permissions: Final[dict[str, list[str]]] = {} - for identifier, tools in mcp_tool_permissions.items(): - if not isinstance(tools, list): - tools = [] - for server_id in sorted(identifier_to_server_ids.get(identifier, [])): - normalized_tool_permissions.setdefault(server_id, []) - normalized_tool_permissions[server_id].extend(tools) - object_permission["mcp_tool_permissions"] = { - server_id: _dedupe_preserving_order(tools) for server_id, tools in normalized_tool_permissions.items() + identifier: _dedupe_preserving_order(tools if isinstance(tools, list) else []) + for identifier, tools in mcp_tool_permissions.items() + if identifier_to_server_ids.get(identifier) } -def _rewrite_object_permission_mcp_identifiers( +def _drop_stale_object_permission_mcp_identifiers( object_permission: ObjectPermissionDict | None, identifier_to_server_ids: dict[str, set[str]], ) -> None: if not object_permission or not isinstance(object_permission, dict): return - _rewrite_object_permission_mcp_servers( + _drop_stale_object_permission_mcp_servers( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) - _rewrite_object_permission_mcp_tool_permissions( + _drop_stale_object_permission_mcp_tool_permissions( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) @@ -615,7 +611,7 @@ async def validate_key_mcp_servers_against_team( "validate_key_mcp_servers_against_team: ignoring stale MCP server identifiers (no longer in registry or DB): %s", sorted(stale_identifiers), ) - _rewrite_object_permission_mcp_identifiers( + _drop_stale_object_permission_mcp_identifiers( object_permission=object_permission, identifier_to_server_ids=identifier_to_server_ids, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1d28ad1914f..78d8ce296b8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -14,7 +14,7 @@ import os import re from collections.abc import Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Any, Final, cast +from typing import TYPE_CHECKING, Annotated, Final, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -69,6 +69,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.utils import LlmProviders +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter @@ -121,7 +122,21 @@ def is_passthrough_request_streaming(request_body: object) -> bool: return bool(request_body.get("stream", False)) -def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, Any]: +def _optional_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _string_keyed_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +async def _json_request_body(request: Request) -> Mapping[str, object]: + return await request.json() + + +def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: """ Build the request metadata carrying key-level spend attribution and the pre-call budget reservation for a router-model passthrough request. @@ -210,7 +225,7 @@ async def llm_passthrough_factory_proxy_route( # anthropic is streaming when 'stream' = True is in the body if request.method == "POST": if "multipart/form-data" not in request.headers.get("content-type", ""): - _request_body = await request.json() + _request_body = await _json_request_body(request) else: _request_body = await get_form_data(request) @@ -383,7 +398,7 @@ async def vllm_proxy_route( endpoint=endpoint, request_query_params=request.query_params, request_headers=_safe_get_request_headers(request), - stream=request_body.get("stream", False), + stream=is_streaming_request, content=None, data=None, files=None, @@ -817,7 +832,7 @@ async def handle_bedrock_passthrough_router_model( # Use the common processing path (same as non-router models) # This ensures all metadata, hooks, and logging are properly initialized - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) data["model"] = model @@ -861,8 +876,8 @@ async def handle_bedrock_count_tokens( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - request_body: dict[str, Any], -) -> dict[str, Any]: + request_body: dict[str, object], +) -> dict[str, object]: """ Handle AWS Bedrock CountTokens API requests. @@ -879,7 +894,7 @@ async def handle_bedrock_count_tokens( handler: Final = BedrockCountTokensHandler() # Extract model from request body - model: Final = request_body.get("model") + model: Final = _optional_str(request_body.get("model")) if not model: raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) @@ -1011,7 +1026,7 @@ async def bedrock_llm_proxy_route( "Bedrock passthrough: Using direct Bedrock model '%s' for endpoint '%s'", model, endpoint ) - data: Final[dict[str, Any]] = {} + data: Final[dict[str, object]] = {} base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) data["method"] = request.method @@ -1110,7 +1125,7 @@ async def bedrock_proxy_route( headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: - data: Final = await request.json() + data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) @@ -1201,7 +1216,7 @@ async def comprehend_medical_proxy_route( ) try: - data: Final = await request.json() + data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail=str(e)) @@ -1412,7 +1427,7 @@ async def assemblyai_proxy_route( is_streaming_request = False # assemblyai is streaming when 'stream' = True is in the body if request.method == "POST": - _request_body: Final = await request.json() + _request_body: Final = await _json_request_body(request) if _request_body.get("stream"): is_streaming_request = True @@ -1519,7 +1534,7 @@ async def azure_proxy_route( endpoint=endpoint, request_query_params=request.query_params, request_headers=_safe_get_request_headers(request), - stream=request_body.get("stream", False), + stream=is_streaming_request, content=None, data=None, files=None, @@ -1606,7 +1621,7 @@ async def azure_proxy_route( extra_headers = auth_credentials.get("headers") or {} - base_target_url = litellm_params.get("api_base") + base_target_url = _optional_str(litellm_params.get("api_base")) if base_target_url is None: raise Exception(f"API base not found for {part}") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( @@ -1727,7 +1742,7 @@ def get_vertex_pass_through_handler( def _override_vertex_params_from_router_credentials( - router_credentials: Any | None, + router_credentials: LiteLLM_ManagedVectorStore | None, vertex_project: str | None, vertex_location: str | None, ) -> tuple[str | None, str | None]: @@ -1747,14 +1762,14 @@ def _override_vertex_params_from_router_credentials( verbose_proxy_logger.debug("Using vector store credentials to override vertex project and location") - litellm_params: Final = router_credentials.get("litellm_params", {}) + litellm_params: Final = _string_keyed_mapping(router_credentials.get("litellm_params")) if not litellm_params: verbose_proxy_logger.warning("Vector store credentials found but litellm_params is empty") return vertex_project, vertex_location # Extract vertex_project and vertex_location from litellm_params - vector_store_project: Final = litellm_params.get("vertex_project") - vector_store_location: Final = litellm_params.get("vertex_location") + vector_store_project: Final = _optional_str(litellm_params.get("vertex_project")) + vector_store_location: Final = _optional_str(litellm_params.get("vertex_location")) if vector_store_project: verbose_proxy_logger.debug( @@ -1762,7 +1777,6 @@ def _override_vertex_params_from_router_credentials( vertex_project, vector_store_project, ) - vertex_project = vector_store_project else: verbose_proxy_logger.warning("Vector store credentials found but missing vertex_project in litellm_params") @@ -1772,11 +1786,10 @@ def _override_vertex_params_from_router_credentials( vertex_location, vector_store_location, ) - vertex_location = vector_store_location else: verbose_proxy_logger.warning("Vector store credentials found but missing vertex_location in litellm_params") - return vertex_project, vertex_location + return vector_store_project or vertex_project, vector_store_location or vertex_location _CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( @@ -1884,8 +1897,8 @@ def _forwarded_headers_for_credentialless_vertex_passthrough( async def _prepare_vertex_auth_headers( request: Request, - vertex_credentials: Any | None, - router_credentials: Any | None, + vertex_credentials: VertexPassThroughCredentials | None, + router_credentials: LiteLLM_ManagedVectorStore | None, vertex_project: str | None, vertex_location: str | None, base_target_url: str | None, @@ -1982,7 +1995,7 @@ async def _base_vertex_proxy_route( fastapi_response: Response, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, user_api_key_dict: UserAPIKeyAuth | None = None, - router_credentials: Any | None = None, + router_credentials: LiteLLM_ManagedVectorStore | None = None, ): """ Base function for Vertex AI passthrough routes. @@ -2152,8 +2165,6 @@ async def vertex_discovery_proxy_route( """ import re - from litellm.types.vector_stores import LiteLLM_ManagedVectorStore - # Extract vector store ID from endpoint if present (e.g., dataStores/test-litellm-app_1761094730750) vector_store_credentials: LiteLLM_ManagedVectorStore | None = None vector_store_id_match: Final = re.search(r"dataStores/([^/]+)", endpoint) @@ -3098,7 +3109,7 @@ async def watsonx_proxy_route( is_streaming_request = False if request.method == "POST": if "multipart/form-data" not in request.headers.get("content-type", ""): - _request_body = await request.json() + _request_body = await _json_request_body(request) else: _request_body = await get_form_data(request) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index ee9a5d94440..49ec18013b5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -267,7 +267,7 @@ class VertexPassthroughLoggingHandler: model: Final = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) - _json_response: Final = httpx_response.json() + _json_response: Final[dict[str, object]] = httpx_response.json() litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse() if vertex_image_generation_class.is_image_generation_response(_json_response): @@ -422,7 +422,7 @@ class VertexPassthroughLoggingHandler: - Creates standard logging object - Logs in litellm callbacks """ - kwargs: dict[str, Any] = {} + kwargs: dict[str, object] = {} vertex_location: Final = get_vertex_location_from_url(url_route) if vertex_location is not None: litellm_logging_obj.optional_params["vertex_location"] = vertex_location diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ddd27f0fc2a..79d5d0a016f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -812,6 +812,8 @@ def _resolve_team_callback_wiring( else { # mutable-ok: Logging arg **callback_vars, TRUSTED_CALLBACK_VARS_FIELD: callback_vars, + "metadata": {}, # mutable-ok: Logging arg + "model_info": {}, # mutable-ok: Logging arg } ) return _TeamCallbackWiring( @@ -3233,6 +3235,14 @@ def _get_pass_through_endpoints_from_config() -> list[PassThroughGenericEndpoint return returned_endpoints +def _config_field_endpoints(response: ConfigFieldInfo) -> list[object] | None: + return response.field_value + + +def _request_app(request: Request) -> FastAPI: + return request.app + + async def _get_pass_through_endpoints_from_db( endpoint_id: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, @@ -3249,7 +3259,7 @@ async def _get_pass_through_endpoints_from_db( except Exception: return [] - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final = _config_field_endpoints(response) if pass_through_endpoint_data is None: return [] @@ -3412,7 +3422,7 @@ async def update_pass_through_endpoints( detail={"error": "No pass-through endpoints found"}, ) - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final[list | None] = _config_field_endpoints(response) if pass_through_endpoint_data is None: raise HTTPException( status_code=404, @@ -3483,7 +3493,7 @@ async def update_pass_through_endpoints( _custom_headers: dict | None = updated_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - route_app: Final[FastAPI] = request.app + route_app: Final = _request_app(request) if updated_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( app=route_app, @@ -3575,7 +3585,7 @@ async def create_pass_through_endpoints( _custom_headers: dict | None = created_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - route_app: Final[FastAPI] = request.app + route_app: Final = _request_app(request) if created_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( app=route_app, @@ -3643,7 +3653,7 @@ async def delete_pass_through_endpoints( response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint - pass_through_endpoint_data: Final[list | None] = response.field_value + pass_through_endpoint_data: Final[list | None] = _config_field_endpoints(response) if response.field_value is None or pass_through_endpoint_data is None: raise HTTPException( status_code=400, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 50cb813c6fa..4be0f556ed7 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,6 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Sequence from typing import Any, Final, Literal import litellm @@ -114,11 +115,7 @@ class PipelineExecutor: # Handle terminal actions if action == "allow": - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": return PipelineExecutionResult( @@ -138,11 +135,7 @@ class PipelineExecutor: # action == "next" → continue to next step # Ran out of steps without a terminal action → default allow - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) @staticmethod async def _run_step( @@ -251,6 +244,45 @@ class PipelineExecutor: return None +def _allow_result( + step_results: Sequence[PipelineStepResult], + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> PipelineExecutionResult: + """Build the terminal-allow result, propagating pipeline modifications without the per-step guardrail override.""" + restored: Final = _restore_request_guardrails(working_data, request_data) + return PipelineExecutionResult( + terminal_action="allow", + step_results=list(step_results), # mutable-ok: PipelineExecutionResult field is a list + modified_data=restored if restored != request_data else None, + ) + + +def _restore_request_guardrails( + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> dict: # mutable-ok: merged back into the request dict, which downstream code mutates + """ + Restore the request's own metadata["guardrails"] activation list. + + _run_step overrides it to [step.guardrail] so should_run_guardrail() allows each + step; letting that override escape via modified_data permanently drops every + independently activated guardrail from later lifecycle stages (post_call, etc.). + """ + working_metadata: Final = working_data.get("metadata") + if not isinstance(working_metadata, dict): + return working_data + request_metadata: Final = request_data.get("metadata") + original_guardrails: Final = request_metadata.get("guardrails") if isinstance(request_metadata, dict) else None + stripped: Final = {k: v for k, v in working_metadata.items() if k != "guardrails"} # mutable-ok: request dict + if original_guardrails is not None: + restored: Final = {**stripped, "guardrails": original_guardrails} # mutable-ok: request dict + return {**working_data, "metadata": restored} # mutable-ok: request dict + if not stripped and not isinstance(request_metadata, dict): + return {k: v for k, v in working_data.items() if k != "metadata"} # mutable-ok: request dict + return {**working_data, "metadata": stripped} # mutable-ok: request dict + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8ac63ba25c9..23932ba7c8c 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1225,6 +1225,7 @@ def run_server( if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None: from litellm.proxy.db.db_url_settings import ( add_missing_query_params, + idle_lifetime_params, reader_shareable_params, unsupported_db_scheme, unsupported_db_scheme_message, @@ -1253,6 +1254,9 @@ def run_server( disable_prepared_statements=db_disable_prepared_statements, extra_params=db_extra_connection_params, ) + lifetime_params: Final = idle_lifetime_params( + general_settings.get("database_max_idle_connection_lifetime") + ) if os.getenv("DATABASE_URL", None) is not None: database_url = get_secret("DATABASE_URL", default_value=None) resolved_url: Final[str | None] = str(database_url) if database_url else None @@ -1270,11 +1274,11 @@ def run_server( writer_url, connection_url_params, ) - os.environ["DATABASE_URL"] = modified_url + os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, lifetime_params) if os.getenv("DIRECT_URL", None) is not None: database_url = os.getenv("DIRECT_URL") modified_url = append_query_params(database_url, connection_url_params) - os.environ["DIRECT_URL"] = modified_url + os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, lifetime_params) # The reader pool is a real pool against the same configured cap, so it # gets the allowlisted pool params. Schema-affecting ones, including any # the operator smuggled in through database_extra_connection_params, stay @@ -1288,10 +1292,13 @@ def run_server( db_lock_timeout, ) os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params( - _with_query_value(read_replica_url, "options", reader_options) - if reader_options - else read_replica_url, - reader_shareable_params(connection_url_params), + add_missing_query_params( + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url, + reader_shareable_params(connection_url_params), + ), + lifetime_params, ) subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index acdd0f99295..672d94077b2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -39,7 +39,7 @@ from typing import ( import anyio import websockets import websockets.exceptions -from pydantic import BaseModel, Json, JsonValue +from pydantic import BaseModel, Json, JsonValue, ValidationError from typing_extensions import NotRequired, ReadOnly, assert_never from litellm._uuid import uuid @@ -253,6 +253,7 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError @@ -263,6 +264,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, @@ -350,7 +352,10 @@ from litellm.proxy.common_utils.load_config_utils import ( get_file_contents_from_s3, ) from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations -from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator +from litellm.proxy.common_utils.model_listing_utils import ( + TeamModelNameTranslator, + configured_display_names, +) from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) @@ -9865,6 +9870,35 @@ class ProxyStartupEvent: replace_existing=True, ) + slack_alerting_args: Final = proxy_logging_obj.slack_alerting_instance.alerting_args + user_spend_check_interval: Final = ( + slack_alerting_args.user_spend_check_interval + if isinstance(slack_alerting_args, SlackAlertingArgs) # pyright: ignore[reportUnnecessaryIsInstance] # tests inject a mock slack_alerting_instance + else SlackAlertingArgs().user_spend_check_interval + ) + + async def _scheduled_user_spend_alerts() -> None: + if ( + await pod_lock_manager.acquire_lock( + cronjob_id=USER_SPEND_ALERTS_JOB_ID, + ttl=max(user_spend_check_interval - 60, 60), + allow_reentrant=False, + ) + is False + ): + return + await proxy_logging_obj.slack_alerting_instance.send_user_spend_alerts() + + scheduler.add_job( + _scheduled_user_spend_alerts, + "interval", + seconds=user_spend_check_interval, + next_run_time=datetime.now(timezone.utc) + timedelta(seconds=10 + random.randint(0, 60)), + id=USER_SPEND_ALERTS_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + if os.getenv("PROMETHEUS_URL"): from zoneinfo import ZoneInfo @@ -10192,7 +10226,8 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): + admin_entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + for response_id, lookup_id in admin_entries: model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -10205,7 +10240,10 @@ async def model_list( if wants_anthropic_format: admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above - return create_anthropic_model_list_response(admin_listing) + return create_anthropic_model_list_response( + admin_listing, + display_names=configured_display_names(admin_entries, llm_router), + ) return dict( data=model_data, @@ -10236,7 +10274,8 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): + entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + for response_id, lookup_id in entries: model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -10249,7 +10288,10 @@ async def model_list( if wants_anthropic_format: listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above - return create_anthropic_model_list_response(listing) + return create_anthropic_model_list_response( + listing, + display_names=configured_display_names(entries, llm_router), + ) return dict( data=model_data, @@ -11061,15 +11103,14 @@ async def audio_speech( if callback_headers: custom_headers.update(callback_headers) - # Determine media type based on model type - media_type = "audio/mpeg" # Default for OpenAI TTS - request_model: Final = data.get("model", "") - if request_model: - request_model_lower: Final = request_model.lower() - if "gemini" in request_model_lower and ( - "tts" in request_model_lower or "preview-tts" in request_model_lower - ): - media_type = "audio/wav" # Gemini TTS returns WAV format after conversion + requested_format: Final = data.get("response_format") + upstream_content_type: Final = ( + response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None + ) + media_type: Final = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=requested_format if isinstance(requested_format, str) else None, + ) return StreamingResponse( _audio_speech_chunk_generator(response), @@ -11085,7 +11126,15 @@ async def audio_speech( ) verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - raise e + if isinstance(e, (ProxyException, HTTPException)): + raise e + raise ProxyException( + message=getattr(e, "message", f"{e}"), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + openai_code=getattr(e, "code", None), + code=getattr(e, "status_code", 500), + ) @router.post( @@ -12473,11 +12522,21 @@ async def supported_openai_params(model: str): --header 'Authorization: Bearer sk-1234' ``` """ + from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider + + global llm_router try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) + resolved_models: Final = llm_router.resolved_litellm_models(model) if llm_router is not None else () + target_model: Final = resolved_models[0] if resolved_models else model + declared_provider: Final = declared_authenticating_provider(target_model) + litellm_model, custom_llm_provider = ( + (target_model.removeprefix(f"{declared_provider}/"), declared_provider) + if declared_provider is not None + else litellm.get_llm_provider(model=target_model)[:2] + ) return { "supported_openai_params": litellm.get_supported_openai_params( - model=model, custom_llm_provider=custom_llm_provider + model=litellm_model, custom_llm_provider=custom_llm_provider ) } except Exception: @@ -14954,17 +15013,25 @@ async def alerting_settings( alerting_args_dict = {} alerting_values = None - allowed_args: Final = { - "slack_alerting": {"type": "Boolean"}, - "daily_report_frequency": {"type": "Integer"}, - "report_check_interval": {"type": "Integer"}, - "budget_alert_ttl": {"type": "Integer"}, - "outage_alert_ttl": {"type": "Integer"}, - "region_outage_alert_ttl": {"type": "Integer"}, - "minor_outage_alert_threshold": {"type": "Integer"}, - "major_outage_alert_threshold": {"type": "Integer"}, - "max_outage_alert_list_size": {"type": "Integer"}, - } + allowed_args: Final = MappingProxyType( + { + "slack_alerting": "Boolean", + "daily_report_frequency": "Integer", + "report_check_interval": "Integer", + "budget_alert_ttl": "Integer", + "outage_alert_ttl": "Integer", + "region_outage_alert_ttl": "Integer", + "minor_outage_alert_threshold": "Integer", + "major_outage_alert_threshold": "Integer", + "max_outage_alert_list_size": "Integer", + "daily_spend_per_user_threshold": "Float", + "monthly_spend_per_user_threshold": "Float", + "spend_anomaly_multiplier": "Float", + "spend_anomaly_baseline_days": "Integer", + "spend_anomaly_min_spend": "Float", + "user_spend_check_interval": "Integer", + } + ) _slack_alerting: Final[SlackAlerting] = proxy_logging_obj.slack_alerting_instance _slack_alerting_args_dict: Final = _slack_alerting.alerting_args.model_dump() @@ -14979,7 +15046,7 @@ async def alerting_settings( _response_obj = ConfigList( field_name="slack_alerting", - field_type=allowed_args["slack_alerting"]["type"], + field_type=allowed_args["slack_alerting"], field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.", field_value=is_slack_enabled, stored_in_db=True if alerting_values is not None else False, @@ -14998,7 +15065,7 @@ async def alerting_settings( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_slack_alerting_args_dict.get(field_name, None), stored_in_db=_stored_in_db, @@ -16426,6 +16493,16 @@ async def update_config_general_settings( detail={"error": f"Invalid type of field value={type(data.field_value)} passed in."}, ) + if data.field_name == "alerting_args": + try: + SlackAlertingArgs.model_validate(data.field_value) + except ValidationError as e: + errors: Final = "; ".join(f"{'.'.join(str(loc) for loc in err['loc'])}: {err['msg']}" for err in e.errors()) + raise HTTPException( + status_code=400, + detail={"error": f"Invalid alerting_args: {errors}"}, + ) + ## get general settings from db db_general_settings: Final = await _config_param_table(prisma_client).find_first( where={"param_name": "general_settings"} diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index a746e9af326..66f8c2ea36f 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -986,6 +986,62 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "QwenCloud", + "provider_display_name": "QwenCloud", + "litellm_provider": "qwencloud", + "credential_fields": [ + { + "key": "api_key", + "label": "QwenCloud API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + "tooltip": "The base URL for QwenCloud. Defaults to https://dashscope-intl.aliyuncs.com/compatible-mode/v1 if not specified.", + "required": true, + "field_type": "text", + "options": null, + "default_value": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" + } + ], + "default_model_placeholder": "gpt-3.5-turbo" + }, + { + "provider": "Qwen_AI_Platform", + "provider_display_name": "Qwen AI Platform", + "litellm_provider": "qwen_ai_platform", + "credential_fields": [ + { + "key": "api_key", + "label": "Qwen AI Platform API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://dashscope.aliyuncs.com/compatible-mode/v1", + "tooltip": "The base URL for Qwen AI Platform. Defaults to https://dashscope.aliyuncs.com/compatible-mode/v1 if not specified.", + "required": true, + "field_type": "text", + "options": null, + "default_value": "https://dashscope.aliyuncs.com/compatible-mode/v1" + } + ], + "default_model_placeholder": "gpt-3.5-turbo" + }, { "provider": "Databricks", "provider_display_name": "Databricks", diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 4d62f1d6d71..db574f859b3 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -7,12 +7,14 @@ Provides: """ import base64 +import json from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse, StreamingResponse +from starlette.datastructures import UploadFile import litellm from litellm._logging import verbose_proxy_logger @@ -45,6 +47,16 @@ if TYPE_CHECKING: router: Final = APIRouter() +def _as_string_keyed_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return value + return None + + +def _response_attr(source: object, name: str) -> object: + return getattr(source, name, None) + + def _raise_vector_store_scan_depth_exceeded() -> None: raise HTTPException( status_code=400, @@ -53,8 +65,8 @@ def _raise_vector_store_scan_depth_exceeded() -> None: def _append_payload_to_scan_stack( - payload_stack: list[tuple[Any, int]], - value: Any, + payload_stack: list[tuple[object, int]], + value: object, next_depth: int, ) -> None: if isinstance(value, dict): @@ -117,7 +129,7 @@ async def _authorize_nested_vector_store_ids( def _build_file_metadata_entry( - response: Any, + response: object, file_data: tuple[str, bytes, str] | None = None, file_url: str | None = None, ) -> Mapping[str, str | int | None]: @@ -135,11 +147,11 @@ def _build_file_metadata_entry( from datetime import datetime, timezone # Extract file_id from response - file_id = None - if hasattr(response, "get"): - file_id = response.get("file_id") - elif hasattr(response, "file_id"): - file_id = response.file_id + mapping_response: Final = _as_string_keyed_mapping(response) + raw_file_id: Final = ( + mapping_response.get("file_id") if mapping_response is not None else _response_attr(response, "file_id") + ) + file_id: Final = raw_file_id if isinstance(raw_file_id, str) else None # Extract file information from file_data tuple filename = None @@ -152,7 +164,7 @@ def _build_file_metadata_entry( content_type = file_data[2] if len(file_data) > 2 else None # Build file metadata entry - file_entry: Final = { + file_entry: Final[dict[str, str | int | None]] = { "file_id": file_id, "filename": filename, "file_url": file_url, @@ -169,7 +181,7 @@ def _build_file_metadata_entry( async def _save_vector_store_to_db_from_rag_ingest( - response: Any, + response: object, ingest_options: Mapping[str, dict[str, str | None]], prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, @@ -197,10 +209,11 @@ async def _save_vector_store_to_db_from_rag_ingest( ) # Handle both dict and object responses - if hasattr(response, "get"): - vector_store_id = response.get("vector_store_id") + mapping_response: Final = _as_string_keyed_mapping(response) + if mapping_response is not None: + vector_store_id = mapping_response.get("vector_store_id") elif hasattr(response, "vector_store_id"): - vector_store_id = response.vector_store_id + vector_store_id = _response_attr(response, "vector_store_id") else: verbose_proxy_logger.warning("Unable to extract vector_store_id from response type: %s", type(response)) return @@ -266,14 +279,13 @@ async def _save_vector_store_to_db_from_rag_ingest( verbose_proxy_logger.info("Vector store %s already exists, appending file to metadata", vector_store_id) # Update existing vector store with new file - existing_metadata = existing_vector_store.vector_store_metadata or {} - if isinstance(existing_metadata, str): - import json + stored_metadata: Final = existing_vector_store.vector_store_metadata or {} + existing_metadata: dict[str, object] = ( + json.loads(stored_metadata) if isinstance(stored_metadata, str) else stored_metadata + ) - existing_metadata = json.loads(existing_metadata) - - ingested_files: Final = existing_metadata.get("ingested_files", []) - ingested_files.append(file_entry) + previous_files: Final = existing_metadata.get("ingested_files", []) + ingested_files: Final = [*previous_files, file_entry] if isinstance(previous_files, list) else [file_entry] existing_metadata["ingested_files"] = ingested_files # Update the vector store @@ -340,9 +352,9 @@ async def parse_rag_ingest_request( # Get file file_obj = form_data.get("file") - if file_obj is not None and hasattr(file_obj, "read"): + if isinstance(file_obj, UploadFile): file_content = await file_obj.read(MAX_UPLOAD_SIZE_BYTES + 1) - file_data = (file_obj.filename, file_content, file_obj.content_type) + file_data = (file_obj.filename or "", file_content, file_obj.content_type or "") # Parse JSON from 'request' form field (contains full request body as JSON) request_json_str: Final[str | bytes | None] = form_data.get("request") diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index aa7595ed13d..5907ffc64eb 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -52,7 +52,7 @@ _TOOL_PAYLOAD_KEYS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( "function": ("name", "description", "parameters", "strict"), } ) -_EMPTY_TOOL_PAYLOAD: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_TOOL_PAYLOAD: Final[Mapping[str, object]] = MappingProxyType({}) def _convert_tool_payload_value(key: str, value: object, *, to_chat: bool) -> object: @@ -105,7 +105,7 @@ def _normalize_tool_dialect( return {**data, **{key: value for key, value in replaceable if key in data}} # mutable-ok: plain body dict -def _is_chat_completions_body(data: Mapping[str, Any]) -> bool: +def _is_chat_completions_body(data: Mapping[str, object]) -> bool: messages: Final = data.get("messages") if isinstance(messages, list) and messages: return True @@ -1373,7 +1373,7 @@ async def _enforce_responses_ws_first_frame_model_auth( request: Request, model: str, user_api_key_dict: UserAPIKeyAuth, - llm_router: Any | None, + llm_router: "Router | None", ) -> None: from litellm.proxy.auth.user_api_key_auth import ( _enforce_key_and_fallback_model_access, @@ -1417,7 +1417,7 @@ async def _enforce_responses_ws_first_frame_model_auth( async def responses_websocket_endpoint( websocket: WebSocket, model: str | None = fastapi.Query(None, description="The model to use for the responses WebSocket session."), - user_api_key_dict=Depends(user_api_key_auth_websocket), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): """ Responses API WebSocket mode endpoint. @@ -1462,7 +1462,7 @@ async def responses_websocket_endpoint( return model, first_message = result - data: dict[str, Any] = { + data: dict[str, object] = { "model": model, "websocket": websocket, } @@ -1471,7 +1471,7 @@ async def responses_websocket_endpoint( # Construct a synthetic Request for pre-call processing headers_list: Final = list(websocket.scope.get("headers") or []) - scope: Final[dict[str, Any]] = { + scope: Final[dict[str, object]] = { "type": "http", "method": "POST", "path": "/v1/responses", diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index b03122966c0..0fd242f2bc1 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -76,17 +76,17 @@ class _StreamEventParser: async def background_streaming_task( polling_id: str, - data: dict, + data: dict[str, object], polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, - general_settings: dict, + general_settings: dict[str, object], llm_router: "Router | None", proxy_config: "ProxyConfig", proxy_logging_obj: "ProxyLogging", - select_data_generator, - user_model, + select_data_generator: Callable[..., object] | None, + user_model: str | None, user_temperature: float | None, user_request_timeout: float | None, user_max_tokens: int | None, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 91a0c68fd58..3d0bd5e61c9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -50,12 +50,12 @@ def _route_user_config_request(data: dict, route_type: str): return ret_val -def _is_a2a_agent_model(model_name: Any) -> bool: +def _is_a2a_agent_model(model_name: object) -> bool: """Check if the model name is for an A2A agent (a2a/ prefix).""" return isinstance(model_name, str) and model_name.startswith("a2a/") -def _raise_if_model_fully_blocked(llm_router: LitellmRouter, model_name: Any, team_id: str | None) -> None: +def _raise_if_model_fully_blocked(llm_router: LitellmRouter, model_name: object, team_id: str | None) -> None: if not isinstance(model_name, str) or not model_name: return if not isinstance(llm_router, litellm.Router): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a9c468c3e9d..b44935ce3dd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1533,7 +1533,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1557,6 +1558,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2d113cfe355..91d2ece7a51 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import Any, Final, NoReturn, cast +from typing import Final, NoReturn, SupportsFloat, SupportsIndex, SupportsInt, cast from fastapi import HTTPException, status @@ -35,6 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget +from litellm.types.router import DeploymentTypedDict @dataclass @@ -119,13 +120,15 @@ async def _apply_over_budget_reservation_policy( applied_entries: list[dict[str, float | str]], reservation_cost: float, current_spend: float, + fail_closed_budget_enforcement: bool = False, ) -> float: """ Decide what to do when a counter is over budget, and return the reservation cost to carry into the next counter. Three outcomes: an over-budget key that opted into throttling releases its own reservation (the rate limiter slows it) and keeps the cost; a partially-remaining budget resizes the reservation - down to what is left; anything else hard-blocks by raising. + down to what is left, unless strict enforcement is on, because the known + estimate already does not fit; anything else hard-blocks by raising. """ if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) @@ -133,21 +136,36 @@ async def _apply_over_budget_reservation_policy( return reservation_cost remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation > 1e-12: - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, + if remaining_before_reservation <= 1e-12: + _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) + if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: + _raise_counter_budget_exceeded( + counter=counter, + current_cost=current_spend - reservation_cost, + estimated_cost=reservation_cost, ) - return remaining_before_reservation + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + return remaining_before_reservation + +def _raise_counter_budget_exceeded( + counter: _BudgetCounter, + current_cost: float, + estimated_cost: float | None = None, +) -> NoReturn: + estimate_detail: Final = "" if estimated_cost is None else f"Estimated request cost: {estimated_cost}, " raise litellm.BudgetExceededError( - current_cost=current_spend, + current_cost=current_cost, max_budget=counter.max_budget, message=( "Budget has been exceeded! " f"{counter.entity_type}={counter.entity_id} " - f"Current cost: {current_spend}, " + f"Current cost: {current_cost}, " + f"{estimate_detail}" f"Max budget: {counter.max_budget}" ), entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type), @@ -257,6 +275,7 @@ async def reserve_budget_for_request( applied_entries=applied_entries, reservation_cost=reservation_cost, current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) continue except Exception: @@ -697,7 +716,7 @@ def _get_budget_limit_counters( for window in budget_limits: window_dict = _coerce_window(window) budget_duration = window_dict.get("budget_duration") - max_budget = window_dict.get("max_budget") + max_budget = _to_float(window_dict.get("max_budget")) if not budget_duration or max_budget is None or max_budget <= 0: continue window_start = get_budget_window_start(window_dict) @@ -724,18 +743,20 @@ def _get_budget_limit_counters( return counters -def _coerce_window(window: Any) -> dict: - if isinstance(window, dict): +def _coerce_window(window: object) -> Mapping[str, object]: + if isinstance(window, Mapping): return window if isinstance(window, str): try: - parsed: Final = json.loads(window) - return parsed if isinstance(parsed, dict) else {} + parsed: Final[object] = json.loads(window) except Exception: return {} - if hasattr(window, "model_dump"): - return window.model_dump() - return {} + return parsed if isinstance(parsed, Mapping) else {} + model_dump: Final = getattr(window, "model_dump", None) + if not callable(model_dump): + return {} + dumped: Final[object] = model_dump() + return dumped if isinstance(dumped, Mapping) else {} async def _reserve_counter( @@ -953,7 +974,7 @@ def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float return default_reserved_cost -def get_budget_window_start(window: Any) -> datetime | None: +def get_budget_window_start(window: object) -> datetime | None: window_dict: Final = _coerce_window(window) budget_duration: Final = window_dict.get("budget_duration") if budget_duration is None: @@ -971,7 +992,7 @@ def get_budget_window_start(window: Any) -> datetime | None: return reset_at - timedelta(seconds=duration_seconds) -def _coerce_datetime(value: Any) -> datetime | None: +def _coerce_datetime(value: object) -> datetime | None: if value is None: return None if isinstance(value, datetime): @@ -1245,11 +1266,11 @@ def _get_model_cost_infos( def _deployment_tiered_pricing_table( - deployment: dict[str, Any], + deployment: DeploymentTypedDict, llm_router: Router, -) -> list[dict] | None: - model_id: Final = deployment.get("model_info", {}).get("id") - backend_model: Final = deployment.get("litellm_params", {}).get("model") +) -> Sequence[Mapping[str, object]] | None: + model_id: Final = _get_value(_get_value(deployment, "model_info"), "id") + backend_model: Final = _get_value(_get_value(deployment, "litellm_params"), "model") if not isinstance(model_id, str) or not isinstance(backend_model, str): return None deployment_model_info: Final = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model) @@ -1414,7 +1435,7 @@ def _estimate_output_tokens( return min(requested, model_ceiling) -def _count_text_tokens(model: str, text: Any) -> int: +def _count_text_tokens(model: str, text: object) -> int: if text is None: return 0 @@ -1454,8 +1475,8 @@ def _is_input_only_route(route: str) -> bool: ) -def _to_float(value: Any) -> float | None: - if value is None: +def _to_float(value: object) -> float | None: + if not isinstance(value, (SupportsFloat, SupportsIndex, str, bytes, bytearray)): return None try: return float(value) @@ -1463,8 +1484,8 @@ def _to_float(value: Any) -> float | None: return None -def _to_int(value: Any) -> int | None: - if value is None: +def _to_int(value: object) -> int | None: + if not isinstance(value, (SupportsInt, SupportsIndex, str, bytes, bytearray)): return None try: return int(value) @@ -1472,7 +1493,7 @@ def _to_int(value: Any) -> int | None: return None -def _get_value(obj: Any, key: str) -> Any: - if isinstance(obj, dict): +def _get_value(obj: object, key: str) -> object: + if isinstance(obj, Mapping): return obj.get(key) return getattr(obj, key, None) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index da5ce2f727e..9f98f1afaa6 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,10 +1,10 @@ import os import re import secrets -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel @@ -222,7 +222,28 @@ def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | return resolved_id -def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> dict: +_MISSING_ATTRIBUTE: Final = object() + + +def _attribute_or_missing(source: object, name: str) -> object: + return getattr(source, name, _MISSING_ATTRIBUTE) + + +@runtime_checkable +class _ModelDumpable(Protocol): + def model_dump(self) -> object: ... + + +def _dumped_usage_info(usage_info: object) -> object: + if isinstance(usage_info, _ModelDumpable): + return usage_info.model_dump() + instance_dict: Final = _attribute_or_missing(usage_info, "__dict__") + if instance_dict is not _MISSING_ATTRIBUTE: + return instance_dict + return usage_info + + +def _extract_usage_for_ocr_call(response_obj: object, response_obj_dict: dict) -> dict: """ Extract usage information for OCR/AOCR calls. @@ -243,12 +264,10 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d usage_info = response_obj_dict.get("usage_info") # Try to extract usage_info from object attributes if not found in dict - if not usage_info and hasattr(response_obj, "usage_info"): - usage_info = response_obj.usage_info - if hasattr(usage_info, "model_dump"): - usage_info = usage_info.model_dump() - elif hasattr(usage_info, "__dict__"): - usage_info = vars(usage_info) + if not usage_info: + attribute_usage_info: Final = _attribute_or_missing(response_obj, "usage_info") + if attribute_usage_info is not _MISSING_ATTRIBUTE: + usage_info = _dumped_usage_info(attribute_usage_info) # For OCR, we track pages instead of tokens if usage_info is not None: @@ -621,6 +640,14 @@ def _ensure_datetime_utc(timestamp: datetime) -> datetime: return timestamp +async def _query_raw_rows( + prisma_client: PrismaClient, + sql_query: str, + *args: object, +) -> Sequence[Mapping[str, object]] | None: + return await prisma_client.db.query_raw(sql_query, *args) + + async def get_spend_by_team( start_date: dt, end_date: dt, @@ -682,7 +709,7 @@ async def get_spend_by_team( group_by_day; """ - db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id) + db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id) if db_response is None: return [] @@ -757,7 +784,7 @@ async def get_spend_by_team_and_customer( group_by_day; """ - db_response: Final = await prisma_client.db.query_raw(sql_query, start_date, end_date, team_id, customer_id) + db_response: Final = await _query_raw_rows(prisma_client, sql_query, start_date, end_date, team_id, customer_id) if db_response is None: return [] @@ -812,7 +839,7 @@ def _sanitize_request_body_for_spend_logs_payload( return {} visited.add(obj_id) - def _sanitize_value(value: Any) -> Any: + def _sanitize_value(value: object) -> object: if isinstance(value, dict): return _sanitize_request_body_for_spend_logs_payload(value, visited, max_string_length_prompt_in_db) elif isinstance(value, list): @@ -1107,7 +1134,7 @@ def _sanitize_error_information_for_spend_logs( return cast(StandardLoggingPayloadErrorInformation, sanitized) -def _convert_to_json_serializable_dict(obj: Any, visited: set | None = None, max_depth: int = 20) -> Any: +def _convert_to_json_serializable_dict(obj: object, visited: set[int] | None = None, max_depth: int = 20) -> object: """ Convert object to JSON-serializable dict, handling Pydantic models safely. @@ -1161,6 +1188,13 @@ def _convert_to_json_serializable_dict(obj: Any, visited: set | None = None, max visited.remove(obj_id) +def _convert_mapping_to_json_serializable(obj: Mapping[str, object]) -> dict[str, object]: + converted: Final = _convert_to_json_serializable_dict(obj) + if isinstance(converted, dict): + return converted + return dict(obj) + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, litellm_params: dict, @@ -1197,7 +1231,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): - _request_body = _convert_to_json_serializable_dict(_request_body) + _request_body = _convert_mapping_to_json_serializable(_request_body) perform_redaction(model_call_details=_request_body, result=None) _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) @@ -1242,7 +1276,7 @@ def _get_response_for_spend_logs_payload( if payload is None: return "{}" if _should_store_prompts_and_responses_in_spend_logs(): - response_obj: Any = payload.get("response") + response_obj: object = payload.get("response") if response_obj is None: return "{}" diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index a1eb7ed06eb..52258602581 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -3,10 +3,11 @@ import asyncio import json import os from collections import Counter -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import ( - Any, Final, + NamedTuple, Protocol, cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read ) @@ -15,6 +16,7 @@ from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -44,6 +46,31 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +JsonSchemaItems: Final = TypedDict( + "JsonSchemaItems", + {"$ref": ReadOnly[str], "enum": ReadOnly[Sequence[JsonValue]]}, + total=False, +) + + +class JsonSchemaNode(TypedDict, total=False): + type: ReadOnly[str] + description: ReadOnly[str] + enum: ReadOnly[Sequence[JsonValue]] + anyOf: ReadOnly[Sequence["JsonSchemaNode"]] + items: ReadOnly["JsonSchemaItems"] + properties: ReadOnly[Mapping[str, "JsonSchemaNode"]] + + +_EMPTY_SCHEMA_DEFS: Final[Mapping[str, "JsonSchemaNode"]] = MappingProxyType({}) + + +class JsonSchemaPropertyEntry(TypedDict): + description: ReadOnly[str] + type: ReadOnly[str] + items: NotRequired[ReadOnly["JsonSchemaItems"]] + + class _SsoSettingsMappingRow(Protocol): @property def sso_settings(self) -> Mapping[str, object] | None: ... @@ -157,10 +184,10 @@ class UIThemeConfig(BaseModel): class SettingsResponse(BaseModel): """Base response model for settings with values and schema information""" - values: dict[str, Any] + values: dict[str, object] """The current configuration values""" - field_schema: dict[str, Any] + field_schema: dict[str, object] """Schema information including descriptions and property types for UI display""" @@ -548,6 +575,62 @@ async def delete_allowed_ip( return {"message": f"IP {ip_address.ip} deleted successfully", "status": "success"} +def _resolve_non_null_variant(field_info: JsonSchemaNode) -> JsonSchemaNode: + """Pydantic v2 renders Optional fields as ``anyOf: [actual_type, null]``.""" + if "anyOf" not in field_info: + return field_info + return next((variant for variant in field_info["anyOf"] if variant.get("type") != "null"), field_info) + + +def _schema_items_entry(resolved: JsonSchemaNode, defs: Mapping[str, JsonSchemaNode]) -> "JsonSchemaItems | None": + """Items info (including enum values) for array fields, so the UI can render a multi-select dropdown.""" + if "items" not in resolved: + return None + items: Final = resolved["items"] + if "$ref" not in items: + return items + ref_def: Final = defs.get(items["$ref"].split("/")[-1]) + if ref_def is None or "enum" not in ref_def: + return None + enum_items: Final[JsonSchemaItems] = {"enum": ref_def["enum"]} + return enum_items + + +def _schema_property_entry(field_info: JsonSchemaNode, defs: Mapping[str, JsonSchemaNode]) -> JsonSchemaPropertyEntry: + resolved: Final = _resolve_non_null_variant(field_info) + items_entry: Final = _schema_items_entry(resolved, defs) + description: Final = field_info.get("description", "") + type_name: Final = resolved.get("type", "string") + if items_entry is None: + entry: Final[JsonSchemaPropertyEntry] = {"description": description, "type": type_name} + return entry + entry_with_items: Final[JsonSchemaPropertyEntry] = { + "description": description, + "type": type_name, + "items": items_entry, + } + return entry_with_items + + +class _RootSchema(NamedTuple): + description: str + properties: Mapping[str, JsonSchemaNode] + nested_defs: Mapping[str, JsonSchemaNode] + defs: Mapping[str, JsonSchemaNode] + + +def _root_schema(settings_class: type[BaseModel]) -> _RootSchema: + from pydantic import TypeAdapter + + raw_schema: Final = TypeAdapter(settings_class).json_schema(by_alias=True) + return _RootSchema( + description=raw_schema.get("description", ""), + properties=raw_schema["properties"], + nested_defs=raw_schema.get("definitions", _EMPTY_SCHEMA_DEFS), + defs=raw_schema["$defs"] if "$defs" in raw_schema else raw_schema.get("definitions", _EMPTY_SCHEMA_DEFS), + ) + + async def _get_settings_with_schema( settings_key: str, settings_class: type[BaseModel], @@ -561,69 +644,43 @@ async def _get_settings_with_schema( settings_class: The Pydantic class to use for schema config: The config dictionary """ - from pydantic import TypeAdapter - litellm_settings: Final = config.get("litellm_settings", {}) or {} settings_data: Final = litellm_settings.get(settings_key, {}) or {} # Create the settings object settings: Final = settings_class(**(settings_data)) # Get the schema - schema: Final = TypeAdapter(settings_class).json_schema(by_alias=True) + root_schema: Final = _root_schema(settings_class) # Convert to dict for response settings_dict: Final = settings.model_dump() # Add descriptions to the response - result: Final = { - "values": settings_dict, - "field_schema": { - "description": schema.get("description", ""), - "properties": {}, - }, + schema_properties_out: Final[Mapping[str, JsonSchemaPropertyEntry]] = { + field_name: _schema_property_entry(field_info, root_schema.defs) + for field_name, field_info in root_schema.properties.items() } - # Add property descriptions - defs: Final = schema.get("$defs", schema.get("definitions", {})) - for field_name, field_info in schema["properties"].items(): - # For Optional fields, Pydantic v2 uses anyOf with [actual_type, null]. - # Resolve the non-null variant to get the real type and items. - resolved = field_info - if "anyOf" in field_info: - for variant in field_info["anyOf"]: - if variant.get("type") != "null": - resolved = variant - break - - prop_entry: dict = { - "description": field_info.get("description", ""), - "type": resolved.get("type", "string"), - } - # Pass through items info (including enum values) for array fields - # so the UI can render a multi-select dropdown - if "items" in resolved: - items = resolved["items"] - # Resolve $ref to enum definitions if needed - if "$ref" in items: - ref_name = items["$ref"].split("/")[-1] - ref_def = defs.get(ref_name, {}) - if "enum" in ref_def: - prop_entry["items"] = {"enum": ref_def["enum"]} - else: - prop_entry["items"] = items - result["field_schema"]["properties"][field_name] = prop_entry - # Add nested object descriptions - for def_name, def_schema in schema.get("definitions", {}).items(): - result["field_schema"][def_name] = { + nested_defs_out: Final[Mapping[str, Mapping[str, object]]] = { + def_name: { "description": def_schema.get("description", ""), "properties": { prop_name: {"description": prop_info.get("description", "")} for prop_name, prop_info in def_schema.get("properties", {}).items() }, } + for def_name, def_schema in root_schema.nested_defs.items() + } - return result + return { + "values": settings_dict, + "field_schema": { + "description": root_schema.description, + "properties": schema_properties_out, + **nested_defs_out, + }, + } @router.get( @@ -930,32 +987,29 @@ async def get_sso_settings(): resolved: Final = resolve_sso_config(sso_db_settings, os.environ) # Get the schema for UI display - from pydantic import TypeAdapter - - schema: Final = TypeAdapter(SSOConfig).json_schema(by_alias=True) + root_schema: Final = _root_schema(SSOConfig) # Convert to dict for response, masking OAuth client secrets so plaintext # is never sent to the UI. sso_dict: Final = mask_sensitive_keys(resolved.config.model_dump(), set(SSO_SECRET_FIELDS)) # Add descriptions to the response - result: Final = { - "values": sso_dict, - "provenance": resolved.provenance, - "field_schema": { - "description": schema.get("description", ""), - "properties": {}, - }, - } - - # Add property descriptions - for field_name, field_info in schema["properties"].items(): - result["field_schema"]["properties"][field_name] = { + schema_properties_out: Final[Mapping[str, Mapping[str, str]]] = { + field_name: { "description": field_info.get("description", ""), "type": field_info.get("type", "string"), } + for field_name, field_info in root_schema.properties.items() + } - return result + return { + "values": sso_dict, + "provenance": resolved.provenance, + "field_schema": { + "description": root_schema.description, + "properties": schema_properties_out, + }, + } @router.patch( @@ -1309,7 +1363,7 @@ UI_SETTINGS_CACHE_KEY: Final = "ui_settings:settings_dict" UI_SETTINGS_CACHE_TTL: Final = 600 # 10 minutes -async def get_ui_settings_cached() -> dict[str, Any]: +async def get_ui_settings_cached() -> dict[str, JsonValue]: """ Return the persisted UI settings dict, using DualCache for reads. diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index d985a546fa7..66071c05b4f 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -1,6 +1,6 @@ #### Video Endpoints ##### -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse @@ -161,7 +161,7 @@ async def video_list( # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params} + data: Final[dict[str, object]] = {"query_params": query_params} # Extract custom_llm_provider from headers, query params, or body custom_llm_provider: Final = ( @@ -246,7 +246,7 @@ async def video_status( ) # Create data with video_id - data: Final[dict[str, Any]] = {"video_id": video_id} + data: Final[dict[str, object]] = {"video_id": video_id} decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") @@ -345,7 +345,7 @@ async def video_content( ) # Create data with video_id - data: Final[dict[str, Any]] = {"video_id": video_id} + data: Final[dict[str, object]] = {"video_id": video_id} decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") @@ -653,7 +653,7 @@ async def video_get_character( ) original_requested_character_id: Final = character_id - data: Final[dict[str, Any]] = {"character_id": character_id} + data: Final[dict[str, object]] = {"character_id": character_id} decoded: Final = decode_character_id_with_provider(character_id) provider_from_id: Final = decoded.get("custom_llm_provider") diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 2dcaa200cc6..7bc1a6a52a3 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -29,6 +29,7 @@ from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.ingestion.vertex_ai_ingestion import VertexAIRAGIngestion from litellm.rag.rag_query import RAGQuery +from litellm.types.llms.openai import AllMessageValues from litellm.types.rag import ( RAGIngestOptions, RAGIngestResponse, @@ -204,7 +205,7 @@ def _suppressed_sub_call_billing() -> Iterator[None]: async def _execute_query_pipeline( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, @@ -311,7 +312,7 @@ async def _execute_query_pipeline( @client async def aquery( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, @@ -358,12 +359,12 @@ async def aquery( @client def query( model: str, - messages: list[Any], + messages: list[AllMessageValues], retrieval_config: dict[str, Any], rerank: dict[str, Any] | None = None, stream: bool = False, **kwargs, -) -> ModelResponse | Coroutine[Any, Any, ModelResponse]: +) -> ModelResponse | Coroutine[None, None, ModelResponse]: """ Query a RAG pipeline. """ @@ -410,7 +411,7 @@ def ingest( file_id: str | None = None, timeout: float | httpx.Timeout | None = None, **kwargs, -) -> RAGIngestResponse | Coroutine[Any, Any, RAGIngestResponse]: +) -> RAGIngestResponse | Coroutine[None, None, RAGIngestResponse]: """ Ingest a document into a vector store. diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 9df1bceac9c..87eb45f262d 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -6,15 +6,34 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final -from litellm.models.user import LiteLLM_UserTable +from pydantic import TypeAdapter + +from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict from litellm.repositories.prisma_protocols import TableActions if TYPE_CHECKING: + from prisma import Prisma from prisma import models as prisma_models _JSON_ENCODED_COLUMNS: Final = frozenset({"metadata", "model_spend", "model_max_budget"}) +_SHADOWING_PLACEHOLDERS_SQL: Final = """ +SELECT p.user_id AS placeholder_user_id, + array_agg(r.user_id ORDER BY r.user_id) AS resolved_user_ids, + p.teams AS team_ids +FROM "LiteLLM_UserTable" p +JOIN "LiteLLM_UserTable" r + ON r.user_id <> p.user_id + AND (r.sso_user_id = p.user_id OR LOWER(r.user_email) = LOWER(p.user_id)) +WHERE p.sso_user_id IS NULL + AND NOT EXISTS (SELECT 1 FROM "LiteLLM_VerificationToken" k WHERE k.user_id = p.user_id) +GROUP BY p.user_id, p.teams +ORDER BY p.user_id +""" + +_PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...]) + class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @@ -59,6 +78,11 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): """Find all users in a team.""" return await self.find_many(where={"teams": {"has": team_id}}) + async def find_shadowing_placeholders(self, tx: "Prisma") -> tuple[SCIMPlaceholder, ...]: + """Users with no SSO id and no virtual keys whose id is another user's SSO id or email.""" + rows: Final = await tx.query_raw(_SHADOWING_PLACEHOLDERS_SQL) + return _PLACEHOLDER_ROWS_ADAPTER.validate_python(rows) + async def count_billable_users(self) -> int: """Number of users that count toward the license seat limit. diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index c8f7842aebf..37ca989b8d3 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -6,6 +6,7 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler @@ -43,10 +44,23 @@ async def arerank( """ Async: Reranks a list of documents based on their relevance to the query """ + _custom_llm_provider: str | None = ( + None # rebind-ok: set by the declared-provider guard or the get_llm_provider unpack; read in the except + ) try: loop: Final = asyncio.get_event_loop() kwargs["arerank"] = True + declared_provider: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared_provider is not None: + _custom_llm_provider = declared_provider # rebind-ok: see pre-declaration above + else: + _, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above + model=model, + custom_llm_provider=custom_llm_provider, + api_base=kwargs.get("api_base", None), + ) + func: Final = partial( rerank, model, @@ -70,7 +84,11 @@ async def arerank( response = init_response return response except Exception as e: - raise e + raise exception_type( + model=model, + custom_llm_provider=_custom_llm_provider or custom_llm_provider, + original_exception=e, + ) @client @@ -115,6 +133,7 @@ def rerank( model_info: Final = kwargs.get("model_info", None) user: Final = kwargs.get("user", None) client: Final = kwargs.get("client", None) + _custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except try: _is_async: Final = kwargs.pop("arerank", False) is True optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -127,7 +146,7 @@ def rerank( ( model, - _custom_llm_provider, + _custom_llm_provider, # rebind-ok: see pre-declaration above dynamic_api_key, dynamic_api_base, ) = litellm.get_llm_provider( @@ -538,4 +557,8 @@ def rerank( return response except Exception as e: verbose_logger.error("Error in rerank: %s", e) - raise exception_type(model=model, custom_llm_provider=custom_llm_provider, original_exception=e) + raise exception_type( + model=model, + custom_llm_provider=_custom_llm_provider or custom_llm_provider, + original_exception=e, + ) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index cccae06c74b..4aa489d9e50 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -17,6 +17,7 @@ logic. import json from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final from pydantic import BaseModel, TypeAdapter, ValidationError @@ -28,6 +29,15 @@ from litellm.types.llms.openai import ( _MAX_ARGUMENTS_LEN: Final = 1_000_000 +TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE: Final = MappingProxyType({"function_call": "fc", "custom_tool_call": "ctc"}) + + +def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str: + prefix: Final = TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE.get(item_type) + if prefix is None or not tool_id or tool_id.startswith(prefix): + return tool_id + return f"{prefix}_{tool_id}" + def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]: """Extract names of tools originally defined as ``type: "custom"``.""" @@ -45,6 +55,21 @@ def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: return tool_name in custom_tool_names +def serialize_tool_call_arguments(raw_arguments: object, default: str = "") -> str: + """Render tool call arguments as the JSON string tool-call schemas require. + + Arguments normally arrive already JSON-encoded, but clients and providers + also send the decoded object. ``str()`` on a dict yields a Python repr with + single quotes, which every downstream JSON parser rejects with errors like + "Expecting ',' delimiter". + """ + if isinstance(raw_arguments, str): + return raw_arguments or default + if raw_arguments is None: + return default + return json.dumps(raw_arguments, default=str) + + def unwrap_custom_tool_arguments(arguments: str) -> str: """Extract the raw content string from JSON-wrapped arguments. @@ -88,7 +113,7 @@ def build_tool_call_item_kwargs( item_type: Final = "custom_tool_call" if custom else "function_call" kwargs: Final[dict[str, str]] = { "type": item_type, - "id": call_id, + "id": openai_shaped_tool_call_item_id(item_type, call_id), "call_id": call_id, "name": name, "status": status, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..bc25f4fffb1 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, + serialize_tool_call_arguments, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -113,6 +114,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} + self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream self._tool_call_id_by_index: dict[int, str] = {} self._ambiguous_tool_call_indexes: set[int] = set() self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item @@ -213,10 +215,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args_delta = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args_delta = str(fn.get("arguments") or "") + fn_args_delta = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args_delta = str(getattr(fn, "arguments", "") or "") + fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) output_index = self._get_or_assign_tool_output_index(call_id) @@ -226,6 +228,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 names = self._custom_tool_names item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + self._tool_item_id_by_call_id[call_id] = item_kwargs["id"] if tool_namespace: item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( @@ -247,7 +250,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, - item_id=call_id, + item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) @@ -284,10 +287,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_args = "" if isinstance(fn, dict): fn_name = str(fn.get("name") or "") - fn_args = str(fn.get("arguments") or "") + fn_args = serialize_tool_call_arguments(fn.get("arguments")) else: fn_name = str(getattr(fn, "name", "") or "") - fn_args = str(getattr(fn, "arguments", "") or "") + fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", "")) tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) # Track if this is a new tool call that wasn't streamed @@ -299,6 +302,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 names = self._custom_tool_names item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + self._tool_item_id_by_call_id[call_id] = item_kwargs["id"] if tool_namespace: item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( @@ -324,7 +328,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 delta_event = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, - item_id=call_id, + item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) @@ -334,7 +338,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 done_event = FunctionCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, - item_id=call_id, + item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, arguments=final_args, ) @@ -344,6 +348,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 names = self._custom_tool_names item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names) + item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"]) if tool_namespace: item_kwargs["namespace"] = tool_namespace item_done_event = OutputItemDoneEvent( @@ -1164,16 +1169,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None: if litellm_model_response: - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None: - usage: Final[object] = getattr(litellm_model_response, "usage", None) - if usage is not None: - setattr( - usage, - "cost", - self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response), - ) - # Transform the response responses_api_response: Final = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4a416f61e1b..5f3e88bb12f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -93,6 +93,8 @@ from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, is_custom_tool_call, + openai_shaped_tool_call_item_id, + serialize_tool_call_arguments, unwrap_custom_tool_arguments, validated_allowed_callers, ) @@ -1010,7 +1012,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], tool_use_type), function=ChatCompletionToolCallFunctionChunk( name=str(function.get("name", "")), - arguments=str(function.get("arguments", "{}")), + arguments=serialize_tool_call_arguments(function.get("arguments"), "{}"), ), index=index, ) @@ -1539,7 +1541,7 @@ class LiteLLMCompletionResponsesConfig: type=cast(Literal["function"], _tool_use_definition.get("type") or "function"), function=ChatCompletionToolCallFunctionChunk( name=function.get("name") or "", - arguments=str(function.get("arguments") or ""), + arguments=serialize_tool_call_arguments(function.get("arguments")), ), index=0, ) @@ -1589,7 +1591,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=ChatCompletionToolCallFunctionChunk( name=f"{namespace}__{raw_name}" if qualify else raw_name, - arguments=str(raw_arguments or ""), + arguments=serialize_tool_call_arguments(raw_arguments), ), index=0, ) @@ -2024,7 +2026,7 @@ class LiteLLMCompletionResponsesConfig: function_definition = tool.function tool_name = function_definition.name or "" tool_id = tool.id or "" - tool_arguments = function_definition.get("arguments") or "" + tool_arguments = serialize_tool_call_arguments(function_definition.get("arguments")) # Check if this is a custom tool if is_custom_tool_call(tool_name, custom_tool_names): @@ -2033,7 +2035,7 @@ class LiteLLMCompletionResponsesConfig: custom_item = CustomToolCallOutputItem( type="custom_tool_call", call_id=tool_id, - id=tool_id, + id=openai_shaped_tool_call_item_id("custom_tool_call", tool_id), name=tool_name, input=input_str, status=function_definition.get("status") or "completed", @@ -2064,7 +2066,7 @@ class LiteLLMCompletionResponsesConfig: name=tool_name, arguments=tool_arguments, call_id=tool_id, - id=tool_id, + id=openai_shaped_tool_call_item_id("function_call", tool_id), type="function_call", status=function_definition.get("status") or "completed", ) @@ -2501,8 +2503,7 @@ class LiteLLMCompletionResponsesConfig: choice=choice, ) message_output_items.extend(image_generation_items) - else: - # Regular message output + elif choice.message.content is not None: message_output_items.append( GenericResponseOutputItem( type="message", @@ -2559,7 +2560,7 @@ class LiteLLMCompletionResponsesConfig: type="function", function=Function( name=tool_call.get("name") or "", - arguments=tool_call.get("arguments") or "", + arguments=serialize_tool_call_arguments(tool_call.get("arguments")), ), ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 197d0c02ba8..367915156d1 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -399,7 +399,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _process_mcp_tools_without_openai_transform( - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth | None", mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], litellm_trace_id: str | None = None, mcp_auth_header: str | None = None, @@ -636,7 +636,7 @@ class LiteLLM_Proxy_MCP_Handler: async def _execute_tool_calls( tool_server_map: dict[str, str], tool_calls: Sequence[object], - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth | None", mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d070f7758fd..7871c85220c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -77,6 +77,16 @@ def _is_json_array(value: object) -> TypeIs[list[object]]: # guard-ok: trivial return isinstance(value, list) +def _optional_str(value: object) -> str | None: + """Keep a JSON payload entry only when it is a string, since the wire format is caller-controlled.""" + return value if isinstance(value, str) else None + + +def _json_array_or_empty(value: object) -> Sequence[object]: + """Narrow a JSON payload entry that the caller iterates, tolerating a missing or malformed value.""" + return value if _is_json_array(value) else () + + def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verifies every value is str return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) @@ -96,10 +106,6 @@ class _GetsLitellmParams(Protocol): def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ... -class _PopsOptionalStr(Protocol): - def __call__(self, key: str, default: None, /) -> str | None: ... - - class _UnmasksPiiText(Protocol): def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... @@ -127,10 +133,6 @@ def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: return fn -def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr: - return fn - - _SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" _UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" @@ -342,7 +344,7 @@ class BaseResponsesAPIStreamingIterator: ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - _item: Final = getattr(openai_responses_api_chunk, "item", None) + _item: Final[object] = getattr(openai_responses_api_chunk, "item", None) if _item is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_item, @@ -350,7 +352,7 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED: - _annotation: Final = getattr(openai_responses_api_chunk, "annotation", None) + _annotation: Final[object] = getattr(openai_responses_api_chunk, "annotation", None) if _annotation is not None: ResponsesAPIRequestUtils._encode_container_id_on_output_item( item=_annotation, @@ -405,23 +407,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[ResponsesAPIResponse | None] = getattr( - openai_responses_api_chunk, "response", None - ) - if response_obj: - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) - if usage_obj is not None: - try: - cost: Final[float | None] = self.logging_obj._response_cost_calculator( - result=response_obj - ) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - # Best-effort usage cost annotation should not break stream replay. - pass + _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() @@ -1021,7 +1007,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: - self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events( + self._events: Sequence[ResponsesAPIStreamingResponse] = build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=self.CHUNK_SIZE, @@ -1088,7 +1074,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: - self._events = _build_synthetic_response_events( + self._events = build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE, @@ -1272,22 +1258,32 @@ def _add_text_like_part_events( ) -def _build_synthetic_response_events( +def _stamp_responses_usage_cost( + response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None +) -> None: + if response_obj is None or logging_obj is None: + return + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + if usage_obj is None: + return + if isinstance(getattr(usage_obj, "cost", None), (int, float)): + return + try: + cost: Final[float | None] = logging_obj._response_cost_calculator(result=response_obj) + except Exception: + return + if isinstance(cost, (int, float)) and cost > 0: + setattr(usage_obj, "cost", cost) + + +def build_synthetic_response_events( *, transformed: ResponsesAPIResponse, - logging_obj: LiteLLMLoggingObj, + logging_obj: LiteLLMLoggingObj | None, chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = _get_openai_response_types() - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Final = transformed.usage if hasattr(transformed, "usage") else None - if usage_obj is not None: - try: - cost: Final[float | None] = logging_obj._response_cost_calculator(result=transformed) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - pass + _stamp_responses_usage_cost(transformed, logging_obj) events: Final[list[ResponsesAPIStreamingResponse]] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), @@ -1310,8 +1306,7 @@ def _build_synthetic_response_events( ) if item_type == "message": - raw_content_parts = output_item_payload.get("content") - content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else [] + content_parts: Sequence[object] = _json_array_or_empty(output_item_payload.get("content")) for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( @@ -1359,9 +1354,8 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - raw_summary_items = output_item_payload.get("summary") - summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else [] - for summary_index, summary in enumerate(summary_items): + summaries: Sequence[object] = _json_array_or_empty(output_item_payload.get("summary")) + for summary_index, summary in enumerate(summaries): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -2518,14 +2512,12 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)("model", None) + requested_model: Final[str | None] = _optional_str(call_kwargs.pop("model", None)) model: Final[str] = ( self.model if requested_model is None or requested_model == self.model_group else requested_model ) - previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)( - "previous_response_id", None - ) + previous_response_id: Final[str | None] = _optional_str(call_kwargs.pop("previous_response_id", None)) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/litellm/router.py b/litellm/router.py index edff8294c3e..3068433b9c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -22,7 +22,7 @@ import traceback import weakref from collections import defaultdict from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -143,7 +143,11 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _check_non_standard_fallback_format, - get_fallback_model_group, + clear_pre_routing_selection, + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) from litellm.router_utils.get_retry_from_policy import ( @@ -821,6 +825,7 @@ class Router: self._zero_cost_cache: dict[str, bool] = {} self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None self._init_routing_groups(None) + self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.model_group_affinity_config = model_group_affinity_config @@ -4918,6 +4923,19 @@ class Router: ) response = await response + if self._should_raise_anthropic_refusal_error( + model=model, + original_generic_function=original_generic_function, + response=response, + kwargs=kwargs, + ): + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, + ) + + refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape + raise safeguard_refusal_error(model=model, stop_details=refusal_details) + self.success_calls[model_name] += 1 verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) @@ -4964,6 +4982,11 @@ class Router: # fallback to the original reference for any non-picklable value. # The original_generic_function is preserved so the per-attempt # helper knows which underlying API to call on fallback. + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -4973,6 +4996,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): return await self._aresponses_streaming_iterator( response=response, @@ -5030,6 +5061,10 @@ class Router: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, parse_anthropic_error_event, + parse_anthropic_refusal_stop_details, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, ) source_iterator: Final = response @@ -5068,13 +5103,35 @@ class Router: continue if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit - error_event = parse_anthropic_error_event(chunk) + # A transport can split one SSE data line across byte chunks, so pre-content + # detection parses the accumulated buffer plus the current chunk, never the + # chunk alone; the buffer is already capped, which bounds this window too. + parse_window = ( # rebind-ok: freshly computed each iteration, never carried over + b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + else chunk + ) + error_event = parse_anthropic_error_event(parse_window) retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over not has_generated_content and error_event is not None and _is_retriable_anthropic_status(error_event[2]) and not _anthropic_stream_error_is_gateway_verdict(chunk) ) + refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over + parse_anthropic_refusal_stop_details(parse_window) + if not has_generated_content and error_event is None + else None + ) + if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs): + refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) + raise MidStreamFallbackError( + message=refusal_error.message, + model=model, + llm_provider="anthropic", + original_exception=refusal_error, + is_pre_first_chunk=True, + ) if not has_generated_content and not retriable_pending_error and error_event is None: buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) continue @@ -5186,8 +5243,13 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. + fallback_trigger: Final[Exception] = ( + e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -5243,6 +5305,11 @@ class Router: # share, leaking primary-deployment metadata into the mid-stream # fallback request. safe_deep_copy avoids deep-copying the full # kwargs (which can hold non-deepcopyable logging handles/clients). + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -5252,6 +5319,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and hasattr(response, "__aiter__"): return await self._aanthropic_messages_streaming_iterator( response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator @@ -6807,6 +6882,9 @@ class Router: original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") + # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier + # behind the router name, and fallbacks are configured per tier, not per router. + lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) fallback_failure_exception_str = "" if disable_fallbacks is True or original_model_group is None: @@ -6851,15 +6929,15 @@ class Router: ] # Get external fallbacks — handle both standard and non-standard formats external_fallback_group: list | None = None - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: if _check_non_standard_fallback_format(fallbacks=fallbacks): # Non-standard formats (e.g. ["claude-3-haiku"] or # [{"model": "...", "messages": [...]}]) are passed through directly external_fallback_group = fallbacks else: - external_fallback_group, generic_idx = get_fallback_model_group( + external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) if external_fallback_group is None and generic_idx is not None: external_fallback_group = fallbacks[generic_idx]["*"] @@ -6917,9 +6995,9 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: context_window_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=context_window_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if context_window_fallback_model_group is None: @@ -6950,9 +7028,9 @@ class Router: elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: content_policy_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=content_policy_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if content_policy_fallback_model_group is None: @@ -6979,14 +7057,14 @@ class Router: if litellm.expose_router_debug_in_errors: e.message += f"\n{error_message}" - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks)) ( fallback_model_group, generic_fallback_idx, - ) = get_fallback_model_group( + ) = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) ## if none, check for generic fallback if fallback_model_group is None and generic_fallback_idx is not None: @@ -6995,12 +7073,12 @@ class Router: if fallback_model_group is None: masked_fallbacks: Final = mask_sensitive_structure(fallbacks) verbose_router_logger.info( - "No fallback model group found for original model_group=%s. Fallbacks=%s", - model_group, + "No fallback model group found for lookup_groups=%s. Fallbacks=%s", + " -> ".join(lookup_groups), masked_fallbacks, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: - original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}" + original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}" raise original_exception input_kwargs.update( @@ -7046,6 +7124,7 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Final[str | None] = kwargs.get("model") + clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): _fallback_metadata_key: Final = _get_router_metadata_variable_name( function_name=getattr(kwargs.get("original_function"), "__name__", None) @@ -7471,6 +7550,24 @@ class Router: break return fallback_model_group + def _get_fallback_model_group_for_lookup_groups( + self, + fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract + lookup_groups: tuple[str, ...], + ) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract + """First lookup group whose exact-key chain resolves (tier first, then requested group).""" + return next( + ( + resolved + for resolved in ( + self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group) + for group in lookup_groups + ) + if resolved is not None + ), + None, + ) + def _get_first_default_fallback(self) -> str | None: """ Returns the first model from the default_fallbacks list, if it exists. @@ -7886,6 +7983,31 @@ class Router: return True return False + def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + """ + Whether a content-policy fallback would resolve for this request, keyed the same way + async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook + selected wins over the requested group. Raising without this returning True would turn + a deliverable response into an error the fallback chain cannot recover from. + """ + content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + if content_policy_fallbacks is not None: + return ( + self._get_fallback_model_group_for_lookup_groups( + fallbacks=content_policy_fallbacks, + lookup_groups=fallback_lookup_groups(kwargs, model_group), + ) + is not None + ) + if self._has_default_fallbacks(): + return True + verbose_router_logger.debug( + "No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s", + model_group, + content_policy_fallbacks, + ) + return False + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. @@ -7898,27 +8020,26 @@ class Router: if response.choices[0].finish_reason != "content_filter": return False - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + return self._has_content_policy_fallback(model, kwargs) - ### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ### - if content_policy_fallbacks is not None: - fallback_model_group = None - for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}] - if list(item.keys())[0] == model: - fallback_model_group = item[model] - break - - if fallback_model_group is not None: - return True - elif self._has_default_fallbacks(): # default fallbacks set - return True - - verbose_router_logger.debug( - "Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s", - model, - content_policy_fallbacks, + def _should_raise_anthropic_refusal_error( + self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any] + ) -> bool: + """ + The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard + refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only + when a content-policy fallback is configured; a plain refusal without stop_details, or + any response with nothing configured, is returned to the client unchanged. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + get_safeguard_refusal_stop_details, ) - return False + + if getattr(original_generic_function, "__name__", "") != "anthropic_messages": + return False + if get_safeguard_refusal_stop_details(response) is None: + return False + return self._has_content_policy_fallback(model, kwargs) def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): _all_deployments: list = [] @@ -8165,6 +8286,52 @@ class Router: if backend_value is not None: model_info[field] = backend_value + @staticmethod + def _inherit_builtin_base_rates_for_off_peak( + model_info: dict, # mutable-ok: cost-map entry filled in place + backend_model: str, + custom_llm_provider: str | None, + ) -> None: + """Fill missing pricing fields on a deployment entry that only sets + ``off_peak_pricing``, from the backend model's built-in cost map entry. + + Cost lookup selects the deployment-scoped entry over the shared backend + entry only when the deployment entry carries a base pricing field, and + ``off_peak_pricing`` is deliberately kept off the shared entry, so a + deployment spelling out only its off-peak schedule would otherwise + never receive the discount. The backend model's entire canonical cost + map entry is copied, field by field, so threshold, tiered, + service-tier, cache, character, and per-second rates as well as + companion billing fields like ``web_search_billing_unit`` and the + regional uplift multipliers all carry over, and peak-hour billing + through the deployment entry matches the shared backend entry exactly. + The raw ``litellm.model_cost`` entry is the copy source rather than + ``get_model_info``'s view of it, since that view synthesizes zero flat + token rates for backends without one and storing those would mark a + tiered-only backend explicitly priced free. Values are deep-copied to + keep the builtin entry isolated. User-specified fields always win; + no-op when any base pricing field is already set or the backend model + has no canonical entry. + """ + if not model_info.get("off_peak_pricing"): + return + if any( + model_info.get(field) is not None + for field in ("input_cost_per_token", "input_cost_per_second", "tiered_pricing") + ): + return + try: + backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model + return + backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "") + if not isinstance(backend_entry, dict): + return + for field, backend_value in backend_entry.items(): + if model_info.get(field) is not None or backend_value is None: + continue + model_info[field] = copy.deepcopy(backend_value) + @staticmethod def _inherit_builtin_tiered_output_rate( model_info: dict, backend_model: str, custom_llm_provider: str | None @@ -8253,6 +8420,11 @@ class Router: if deployment.litellm_params.get(field) is not None: _model_info[field] = deployment.litellm_params[field] + Router._inherit_builtin_base_rates_for_off_peak( + model_info=_model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) if _model_info.get("input_cost_per_token") is not None: Router._inherit_builtin_cache_pricing( model_info=_model_info, @@ -8301,6 +8473,19 @@ class Router: return deployment except Exception as e: if self.ignore_invalid_deployments: + if isinstance(e, litellm.BadRequestError): + self._provider_unresolved_deployments = ( + *self._provider_unresolved_deployments, + partial( + self._create_deployment, + deployment_info=deployment_info, + _model_name=_model_name, + _litellm_params=_litellm_params, + _model_info=_model_info, + declared_id=declared_id, + duplicate_ids=duplicate_ids, + ), + ) verbose_router_logger.exception( "Error creating deployment: %s, ignoring and continuing with other deployments.", e ) @@ -8730,6 +8915,7 @@ class Router: self.quality_routers = {} self.complexity_routers = {} self.auto_routers = {} + self._provider_unresolved_deployments = () self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -8994,6 +9180,11 @@ class Router: if field_value is not None: _model_info_dict[field] = field_value + Router._inherit_builtin_base_rates_for_off_peak( + model_info=_model_info_dict, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) if _model_info_dict.get("input_cost_per_token") is not None: Router._inherit_builtin_cache_pricing( model_info=_model_info_dict, @@ -9249,6 +9440,11 @@ class Router: field_value = deployment.litellm_params.get(field) if field_value is not None: model_info[field] = field_value + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) if model_info.get("input_cost_per_token") is not None: Router._inherit_builtin_cache_pricing( model_info=model_info, @@ -9342,8 +9538,12 @@ class Router: """Re-assert this router's deployments onto a freshly fetched catalog. Reads ``model_list`` at call time, so only deployments the router still - serves are restored. + serves are restored, plus any config deployment the fresh catalog now resolves. """ + provider_unresolved: Final = self._provider_unresolved_deployments + self._provider_unresolved_deployments = () + for create_deployment in provider_unresolved: + create_deployment() for entry in tuple(self.model_list): try: deployment = entry if isinstance(entry, Deployment) else Deployment(**entry) @@ -9546,6 +9746,26 @@ class Router: coerce_token_limit(model_info.get("max_output_tokens")), ) + def get_configured_display_name(self, model_name: str) -> "str | None": + """ + Return the display_name explicitly configured in a concrete deployment's + model_info for model_name, via O(1) index lookup. + + Returns None for wildcard-expanded or unknown names, and treats a + non-string or empty configured value as absent rather than failing the + listing. Like get_configured_token_limits, this never triggers pattern + matching or deep copies, so it is safe to call per listed model on the + /v1/models hot path. + """ + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + display_name: Final = deployment.model_info.get("display_name") + if isinstance(display_name, str) and display_name.strip(): + return display_name + return None + def get_deployment_credentials_with_provider( self, model_id: str, team_id: str | None = None ) -> dict[str, Any] | None: @@ -12026,6 +12246,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs @@ -12141,6 +12362,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index d57d7da0410..a8d51f95e45 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -20,6 +20,7 @@ anthropic: import asyncio import builtins +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Any, Final @@ -54,19 +55,19 @@ class _LiteLLMParamsDictView: __slots__ = ("_params",) - def __init__(self, params: dict[str, Any]): + def __init__(self, params: Mapping[str, object]): self._params = params - def __getattr__(self, key: str) -> Any: + def __getattr__(self, key: str) -> object: return self._params.get(key) - def __getitem__(self, key: str) -> Any: + def __getitem__(self, key: str) -> object: return self._params.get(key) def __contains__(self, key: str) -> bool: return key in self._params - def get(self, key: str, default: Any = None) -> Any: + def get(self, key: str, default: object = None) -> object: return self._params.get(key, default) def keys(self): @@ -84,10 +85,10 @@ class _LiteLLMParamsDictView: def __len__(self) -> int: return len(self._params) - def dict(self) -> dict[str, Any]: + def dict(self) -> builtins.dict[str, object]: return dict(self._params) - def model_dump(self) -> builtins.dict[str, Any]: + def model_dump(self) -> builtins.dict[str, object]: return dict(self._params) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index be7653902a7..577cee0920d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -282,7 +282,7 @@ def _response_cost_or_none(response: ModelResponse) -> float | None: return float(cost) -def _effective_turn_off_message_logging(request_kwargs: Mapping[str, Any] | None) -> bool | None: +def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None: from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( initialize_standard_callback_dynamic_params, ) @@ -1925,7 +1925,7 @@ class ComplexityRouter(CustomLogger): ceiling_severity: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None best_model: str | None = None best_score = float("-inf") - candidate_scores: Final[list[dict[str, Any]]] = [] + candidate_scores: Final[list[dict[str, object]]] = [] for model in candidates: if floor_severity is not None and all( self._active_tier_severity(model_tier) < floor_severity diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 924574537f3..3d37ca216a7 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: return False +PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model" +_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata") + + +def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None: + """ + Remember which model a pre-routing hook picked, so fallback lookup can key off it. + + Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so + writing the model there is invisible by the time routing picks a tier. The metadata + buckets are nested dicts shared by reference across those copies, which is how the + router already carries values back up. + + The write goes through the proxy-internal bucket resolver, never into both buckets: + on /v1/messages the top-level ``metadata`` dict is the provider's own request field, + so a blanket write would forward the tier stamp upstream. + """ + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if request_kwargs is None: + return + bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) + if isinstance(bucket, dict): + bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model + + +def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None: + """ + Drop any selection the router did not make itself on this hop. + + The buckets carry whatever the caller sent, so an inbound value is the caller + choosing a fallback chain rather than the router choosing a tier. A fallback hop + also inherits the previous hop's selection, which would key its own failure off + the tier that already failed. Clearing at the start of every hop leaves only a + value the pre-routing hook wrote while routing that hop. + """ + if request_kwargs is None: + return + for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS): + if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket: + del bucket[PRE_ROUTING_SELECTED_MODEL_KEY] + + +def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: + """The model a pre-routing hook selected for this request, if one did.""" + buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS) + selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict)) + return next((selected for selected in selections if isinstance(selected, str) and selected), None) + + +def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: + """ + Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, + and the requested group still resolves when no tier-keyed chain exists, so configs keyed + on the router name (the documented contract) keep working behind auto-routers. + """ + ordered: Final = (get_pre_routing_selection(kwargs), model_group) + return tuple(dict.fromkeys(group for group in ordered if group)) + + +def _resolved_a_specific_chain( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract +) -> bool: + resolved, generic_idx = result + if resolved is None: + return False + return generic_idx is None or resolved is not fallbacks[generic_idx]["*"] + + +def get_fallback_model_group_for_lookup_groups( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + lookup_groups: tuple[str, ...], +) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract + """ + First lookup group with a specifically-keyed chain wins; the generic "*" chain applies + only after every group missed, so a catch-all cannot shadow a later group's own chain. + """ + results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups) + specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None) + if specific is not None: + return specific + return next((result for result in results if result[0] is not None), (None, None)) + + def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]: """ Returns: @@ -412,6 +497,7 @@ async def run_async_fallback( # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(mg)) + kwargs.pop("_target_order", None) # rebind-ok: next hop must not inherit the previous order target if isinstance(mg, str): kwargs["model"] = mg elif isinstance(mg, dict): diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 0d5ef01bc04..0775e0a4039 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -8,7 +8,7 @@ from re import Match from typing import Final from litellm._logging import verbose_router_logger -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider class PatternUtils: @@ -204,7 +204,7 @@ class PatternMatchRouter: return litellm_deployment_litellm_model - def get_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict] | None: + def get_pattern(self, model: str | None, custom_llm_provider: str | None = None) -> list[dict] | None: """ Check if a pattern exists for the given model and custom llm provider @@ -215,18 +215,17 @@ class PatternMatchRouter: Returns: bool: True if pattern exists, False otherwise """ - if custom_llm_provider is None: - try: - ( - _, - custom_llm_provider, - _, - _, - ) = get_llm_provider(model=model) - except Exception: - # get_llm_provider raises exception when provider is unknown - pass - return self.route(model) or self.route(f"{custom_llm_provider}/{model}") + provider: Final = ( + custom_llm_provider or declared_authenticating_provider(model) or self._resolved_provider(model) + ) + return self.route(model) or self.route(f"{provider}/{model}") + + @staticmethod + def _resolved_provider(model: str | None) -> str | None: + try: + return get_llm_provider(model=model)[1] if model else None + except Exception: # noqa: BLE001 # get_llm_provider raises when the provider is unknown; the name then routes as-is + return None def get_deployments_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict]: """ diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 7fb90ab89de..b1e9dbdefa8 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -427,6 +427,8 @@ class DeploymentAffinityCheck(CustomLogger): """ request_kwargs = request_kwargs or {} typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) + if request_kwargs.get("_target_order") is not None: + return typed_healthy_deployments ( enable_user_key, diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 6e8406b2ec7..0788c8db710 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -58,6 +58,9 @@ class PromptCachingDeploymentCheck(CustomLogger): request_kwargs: dict | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: + if request_kwargs is not None and request_kwargs.get("_target_order") is not None: + return healthy_deployments + if messages is not None and is_prompt_caching_valid_prompt( messages=messages, model=model, diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index d96defbbcd6..ab5ef5853c9 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -9,6 +9,7 @@ import random import traceback from collections.abc import Callable from functools import partial +from types import MappingProxyType from typing import Any, Final from litellm._logging import verbose_router_logger @@ -214,6 +215,15 @@ class SearchAPIRouter: api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( tool_litellm_params=litellm_params, ) + protected_params: Final = frozenset(("search_provider", "api_key", "api_base")) + search_params: Final = MappingProxyType( + { + key: value + for params in (litellm_params, kwargs) + for key, value in params.items() + if key not in protected_params and value is not None + } + ) verbose_router_logger.debug("Selected search tool with provider: %s", search_provider) @@ -222,7 +232,7 @@ class SearchAPIRouter: search_provider=search_provider, api_key=api_key, api_base=api_base, - **kwargs, + **search_params, ) return response diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 84461115e8e..21f27075e0f 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,16 +2,37 @@ Cost calculation for search providers. """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.utils import get_model_info +PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) +EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _provider_usage( + optional_params: Mapping[str, object] | None, + usage_param: str, +) -> tuple[Mapping[str, object], ...] | None: + params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS + raw_usage: Final[object] = params.get(usage_param) + try: + return PROVIDER_USAGE_ADAPTER.validate_python(raw_usage) + except ValidationError: + return None + def search_provider_cost_per_query( model: str, custom_llm_provider: str | None = None, number_of_queries: int = 1, - optional_params: dict | None = None, + optional_params: Mapping[str, object] | None = None, ) -> tuple[float, float]: """ Calculate cost for search-only providers. @@ -28,6 +49,18 @@ def search_provider_cost_per_query( Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ + if custom_llm_provider == "parallel_ai": + from litellm.llms.parallel_ai.search.cost_calculator import ( + PARALLEL_AI_USAGE_PARAM, + parallel_ai_search_cost, + ) + + input_cost: Final = parallel_ai_search_cost( + optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS, + usage=_provider_usage(optional_params, PARALLEL_AI_USAGE_PARAM), + ) + return (input_cost, 0.0) + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) # Check for tiered pricing (e.g., Exa AI based on max_results) diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index e2662d96b52..8f677b54700 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -1,7 +1,9 @@ import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -17,6 +19,72 @@ from litellm.proxy._types import KeyManagementSystem from .base_secret_manager import BaseSecretManager, raise_if_unsafe_secret_name +class _VaultAuthData(TypedDict): + """The ``auth`` block Vault returns from a login endpoint.""" + + client_token: ReadOnly[str] + lease_duration: ReadOnly[int] + + +class _VaultLoginResponse(TypedDict): + """Body of a Vault ``/v1/auth/.../login`` response.""" + + auth: ReadOnly[_VaultAuthData] + + +class _VaultSecretTarget(TypedDict): + """Resolved coordinates of one Vault KV v2 secret.""" + + url: ReadOnly[str] + data_key: ReadOnly[str] + secret_name: ReadOnly[str] + + +class _VaultSecretDataBlock(TypedDict, total=False): + """The inner ``data`` block of a Vault KV v2 read body.""" + + data: ReadOnly[Mapping[str, object]] + + +class _VaultSecretReadResponse(TypedDict, total=False): + """Body of a Vault KV v2 secret read, narrowed to the nesting this module walks.""" + + data: ReadOnly[_VaultSecretDataBlock] + + +class _VaultLoginResponseSource(Protocol): + """A Vault login call's HTTP response, read for the auth block it carries.""" + + def json(self) -> _VaultLoginResponse: ... + + +class _VaultSecretReadSource(Protocol): + """A Vault KV v2 read response, read for the nested secret data it carries.""" + + def json(self) -> _VaultSecretReadResponse: ... + + +class _JsonObjectSource(Protocol): + """A Vault response whose body is a JSON object nothing further is assumed about.""" + + def json(self) -> dict[str, object]: ... + + +def _vault_login_body(response: _VaultLoginResponseSource) -> _VaultLoginResponse: + """Decode the body of a Vault login response.""" + return response.json() + + +def _vault_secret_read_body(response: _VaultSecretReadSource) -> _VaultSecretReadResponse: + """Decode the body of a Vault KV v2 secret read response.""" + return response.json() + + +def _json_object_body(response: _JsonObjectSource) -> dict[str, object]: + """Decode a Vault response body as a plain JSON object.""" + return response.json() + + class HashicorpSecretManager(BaseSecretManager): def __init__(self): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user @@ -130,7 +198,8 @@ class HashicorpSecretManager(BaseSecretManager): ) resp.raise_for_status() - auth_data: Final = resp.json()["auth"] + login_response: Final = _vault_login_body(resp) + auth_data: Final = login_response["auth"] token: Final = auth_data["client_token"] _lease_duration: Final = auth_data["lease_duration"] @@ -191,8 +260,10 @@ class HashicorpSecretManager(BaseSecretManager): json=self._get_tls_cert_auth_body(), ) resp.raise_for_status() - token: Final = resp.json()["auth"]["client_token"] - _lease_duration: Final = resp.json()["auth"]["lease_duration"] + token_response: Final = _vault_login_body(resp) + token: Final = token_response["auth"]["client_token"] + lease_response: Final = _vault_login_body(resp) + _lease_duration: Final = lease_response["auth"]["lease_duration"] verbose_logger.debug("Successfully obtained Vault token via TLS cert auth.") self.cache.set_cache(key="hcp_vault_token", value=token, ttl=_lease_duration) return token @@ -205,9 +276,9 @@ class HashicorpSecretManager(BaseSecretManager): def get_url( self, secret_name: str, - namespace: str | None = None, - mount_name: str | None = None, - path_prefix: str | None = None, + namespace: object = None, + mount_name: object = None, + path_prefix: object = None, ) -> str: """ Constructs the Vault URL for KV v2 secrets. @@ -238,7 +309,7 @@ class HashicorpSecretManager(BaseSecretManager): _url += secret_name return _url - def _sanitize_plain_value(self, value: str | int | None) -> str | None: + def _sanitize_plain_value(self, value: object) -> str | None: if value is None: return None value_str: Final = str(value).strip() @@ -246,23 +317,23 @@ class HashicorpSecretManager(BaseSecretManager): return None return value_str - def _sanitize_path_component(self, value: str | int | None) -> str | None: + def _sanitize_path_component(self, value: object) -> str | None: sanitized_value = self._sanitize_plain_value(value) if sanitized_value is None: return None sanitized_value = sanitized_value.strip("/") return sanitized_value or None - def _extract_secret_manager_settings(self, optional_params: dict | None) -> dict[str, Any]: + def _extract_secret_manager_settings(self, optional_params: dict | None) -> dict[str, object]: if not isinstance(optional_params, dict): return {} candidate: Final = optional_params.get("secret_manager_settings") - source: Final = candidate if isinstance(candidate, dict) else optional_params + source: Final[Mapping[str, object]] = candidate if isinstance(candidate, dict) else optional_params allowed_keys: Final = {"namespace", "mount", "path_prefix", "data"} return {k: source[k] for k in allowed_keys if k in source} - def _build_secret_target(self, secret_name: str, optional_params: dict | None) -> dict[str, Any]: + def _build_secret_target(self, secret_name: str, optional_params: dict | None) -> _VaultSecretTarget: settings: Final = self._extract_secret_manager_settings(optional_params) namespace: Final = settings.get("namespace", self.vault_namespace) @@ -331,7 +402,7 @@ class HashicorpSecretManager(BaseSecretManager): response.raise_for_status() # For KV v2, the secret is in response.json()["data"]["data"] - json_resp: Final = response.json() + json_resp: Final = _json_object_body(response) _value: Final = self._get_secret_value_from_json_response(json_resp) self.cache.set_cache(secret_name, _value) return _value @@ -362,7 +433,7 @@ class HashicorpSecretManager(BaseSecretManager): response.raise_for_status() # For KV v2, the secret is in response.json()["data"]["data"] - json_resp: Final = response.json() + json_resp: Final = _json_object_body(response) _value: Final = self._get_secret_value_from_json_response(json_resp) self.cache.set_cache(secret_name, _value) return _value @@ -379,7 +450,7 @@ class HashicorpSecretManager(BaseSecretManager): optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None, tags: dict | list | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Writes a secret to Vault KV v2 using an async HTTPX client. @@ -413,7 +484,7 @@ class HashicorpSecretManager(BaseSecretManager): json=data, ) response.raise_for_status() - return response.json() + return _json_object_body(response) except Exception as e: verbose_logger.exception("Error writing secret to Hashicorp Vault: %s", e) return {"status": "error", "message": str(e)} @@ -500,7 +571,7 @@ class HashicorpSecretManager(BaseSecretManager): headers=self._get_request_headers(), ) response.raise_for_status() - json_resp: Final = response.json() + json_resp: Final = _vault_secret_read_body(response) # Use data_key from target to get the correct value data_key: Final = new_target["data_key"] new_secret_value_from_vault: Final = json_resp.get("data", {}).get("data", {}).get(data_key, None) diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index d6b3dfa3285..dd147aaccee 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -53,11 +53,12 @@ PROVIDERS: Final[list[dict]] = [ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Fable 5, Opus 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 5, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5.1, Fable 5, Opus 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 5, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ + "claude-fable-5-1", "claude-fable-5", "claude-opus-5", "claude-sonnet-5", diff --git a/litellm/types/guardrail_base_init.py b/litellm/types/guardrail_base_init.py new file mode 100644 index 00000000000..9174e8d840f --- /dev/null +++ b/litellm/types/guardrail_base_init.py @@ -0,0 +1,24 @@ +"""Typed view of the scalar keyword payload guardrails forward to ``CustomGuardrail.__init__``. + +Guardrail subclasses collect their base-class options in ``**kwargs`` and splat them into +``super().__init__``. Declaring the payload's shape here lets the checker resolve each +forwarded argument to its real parameter type instead of ``Any``. +""" + +from typing_extensions import ReadOnly, TypedDict + + +class GuardrailBaseInitKwargs(TypedDict, total=False): + guardrail_name: ReadOnly[str | None] + default_on: ReadOnly[bool] + mask_request_content: ReadOnly[bool] + mask_response_content: ReadOnly[bool] + violation_message_template: ReadOnly[str | None] + end_session_after_n_fails: ReadOnly[int | None] + on_violation: ReadOnly[str | None] + realtime_violation_message: ReadOnly[str | None] + on_sensitive_data: ReadOnly[str | None] + sensitive_data_route_to_model: ReadOnly[str | None] + sticky_session_routing: ReadOnly[bool] + run_in_parallel: ReadOnly[bool] + only_scan_new_messages: ReadOnly[bool] diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 4cf4fa62eff..c17103da890 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -136,6 +136,7 @@ class SupportedGuardrailIntegrations(Enum): HEADROOM = "headroom" COMPRESR = "compresr" STRAIKER = "straiker" + ALICE = "alice" class Role(Enum): diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index 7853dda1213..bae876dfdd9 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -4,21 +4,58 @@ Payloads for Datadog LLM Observability Service (LLMObs) API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards """ +from collections.abc import Sequence from typing import Any, Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams +class ToolCall(TypedDict, total=False): + """A tool call on a message, as LLM Obs names its fields.""" + + name: ReadOnly[str] + arguments: ReadOnly[dict[str, Any] | str] # parsed object, or the raw string when it will not parse to one + tool_id: ReadOnly[str] + type: ReadOnly[str] + + +class ToolResult(TypedDict, total=False): + """The result of a tool call, as LLM Obs names its fields.""" + + name: ReadOnly[str] + result: ReadOnly[str] + tool_id: ReadOnly[str] + type: ReadOnly[str] + + +class ToolDefinition(TypedDict, total=False): + """A tool the model was offered on the request.""" + + name: ReadOnly[str] + description: ReadOnly[str] + schema: ReadOnly[dict[str, Any]] + + +class Message(TypedDict, total=False): + """A message on a span, as LLM Obs names its fields.""" + + content: ReadOnly[str] + role: ReadOnly[str] + reasoning_content: ReadOnly[str] + tool_calls: ReadOnly[Sequence[ToolCall]] + tool_results: ReadOnly[Sequence[ToolResult]] + + class InputMeta(TypedDict): - messages: list[ - dict[str, Any] # changed to fit with tool calls + messages: Sequence[ + Message | dict[str, Any] # changed to fit with tool calls ] # Relevant Issue: https://github.com/BerriAI/litellm/issues/9494 class OutputMeta(TypedDict): - messages: list[Any] + messages: Sequence[Any] class DDLLMObsError(TypedDict, total=False): @@ -36,6 +73,7 @@ class Meta(TypedDict, total=False): output: OutputMeta # The span's output information. metadata: dict[str, Any] error: DDLLMObsError | None # Error information on the span + tool_definitions: ReadOnly[Sequence[ToolDefinition]] # The tools offered to the model on this request class LLMMetrics(TypedDict, total=False): @@ -45,6 +83,9 @@ class LLMMetrics(TypedDict, total=False): time_to_first_token: float time_per_output_token: float total_cost: float + cache_read_input_tokens: ReadOnly[float] + cache_write_input_tokens: ReadOnly[float] + non_cached_input_tokens: ReadOnly[float] class LLMObsPayload(TypedDict, total=False): diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 01ed8b08571..8498b6f6d00 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -270,6 +270,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_deployment_rpm_limit", "litellm_remaining_api_key_requests_for_model", "litellm_remaining_api_key_tokens_for_model", + "litellm_api_key_rate_limit_allowed_metric", + "litellm_api_key_rate_limit_used_metric", + "litellm_team_rate_limit_allowed_metric", + "litellm_team_rate_limit_used_metric", "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", @@ -775,6 +779,22 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.MODEL_ID.value, ] + litellm_api_key_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = ( + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ) + + litellm_api_key_rate_limit_used_metric = litellm_api_key_rate_limit_allowed_metric + + litellm_team_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = ( + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ) + + litellm_team_rate_limit_used_metric = litellm_team_rate_limit_allowed_metric + litellm_llm_api_failed_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index b1b7bc3541a..64c0c530e9b 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -91,6 +91,40 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase): default=False, description="If true, the alerting payload will be printed to the console.", ) + daily_spend_per_user_threshold: float | None = Field( + default=None, + gt=0, + allow_inf_nan=False, + description="Alert when a user's spend for the current day (UTC) crosses this USD amount. Off by default.", + ) + monthly_spend_per_user_threshold: float | None = Field( + default=None, + gt=0, + allow_inf_nan=False, + description="Alert when a user's spend for the current calendar month (UTC) crosses this USD amount. Off by default.", + ) + spend_anomaly_multiplier: float = Field( + default=3.0, + gt=0, + allow_inf_nan=False, + description="Flag a user's spend as anomalous when today's spend exceeds this multiple of their trailing daily average.", + ) + spend_anomaly_baseline_days: int = Field( + default=7, + ge=1, + description="Number of trailing days used to compute a user's daily average spend for anomaly detection.", + ) + spend_anomaly_min_spend: float = Field( + default=10.0, + gt=0, + allow_inf_nan=False, + description="Minimum spend (USD) a user must reach today before an anomaly alert can fire. Reduces false positives.", + ) + user_spend_check_interval: int = Field( + default=3600, + ge=60, + description="How often (in seconds) to check per-user spend thresholds and anomalies. Default is hourly.", + ) class DeploymentMetrics(LiteLLMPydanticObjectBase): @@ -138,6 +172,8 @@ class AlertType(str, Enum): budget_alerts = "budget_alerts" spend_reports = "spend_reports" failed_tracking_spend = "failed_tracking_spend" + user_spend_thresholds = "user_spend_thresholds" + user_spend_anomalies = "user_spend_anomalies" # Database alerts db_exceptions = "db_exceptions" @@ -182,6 +218,7 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [ AlertType.budget_alerts, AlertType.spend_reports, AlertType.failed_tracking_spend, + AlertType.user_spend_thresholds, # Database alerts AlertType.db_exceptions, # Report alerts diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 42ca3fd6d4b..4fe1dafc73b 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False): server_tool_use: NotRequired[ReadOnly[ServerToolUsage]] +class AnthropicStopDetails(TypedDict, total=False): + """ + Safeguard verdict accompanying a `stop_reason: "refusal"` response: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback + """ + + category: ReadOnly[str | None] + explanation: ReadOnly[str | None] + + class AnthropicMessagesResponse(TypedDict, total=False): """ Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages @@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False): id: str model: str | None # This represents the Model type from Anthropic role: Literal["assistant"] | None - stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None + stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None + stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]] stop_sequence: str | None type: Literal["message"] | None usage: AnthropicUsage | None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index fcade835cce..32d88da0085 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -327,6 +327,10 @@ class BatchGuardrailReport(BaseModel): """Every record that was redacted or dropped, in file order.""" +_JsonValue: TypeAlias = object +"""Alias for ``object``, usable inside model bodies that declare a field named ``object``.""" + + BATCH_GUARDRAIL_RESPONSE_FIELD: Final = "litellm_batch_guardrail" @@ -1191,7 +1195,7 @@ class ShellToolParam(TypedDict, total=False): type: Required[Literal["shell"] | str] """The type of tool. Use ``\"shell\"``.""" - environment: Required[dict[str, Any]] + environment: Required[dict[str, object]] """Environment config: ``type`` (e.g. ``\"container_auto\"``, ``\"container_reference\"``, ``\"local\"``), optional ``container_id``, ``network_policy``, ``domain_secrets``, ``skills``.""" @@ -1308,7 +1312,7 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): @field_validator("cost", mode="before") @classmethod - def parse_cost(cls, v: Any) -> float | None: + def parse_cost(cls, v: object) -> object: """Normalise cost: accept either a float or a dict with a ``total_cost`` key.""" if isinstance(v, dict): return v.get("total_cost") @@ -1805,7 +1809,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): type: str # e.g., 'invalid_request_error' code: str # e.g., 'context_length_exceeded' message: str - param: str | dict[str, Any] | None = None + param: str | dict[str, object] | None = None class ErrorEvent(BaseLiteLLMOpenAIResponseObject): @@ -2418,7 +2422,7 @@ class OpenAIVideoObject(BaseModel): expires_at: int | None = None """Unix timestamp (seconds) for when the downloadable assets expire, if set.""" - error: dict[str, Any] | None = None + error: dict[str, _JsonValue] | None = None """Error payload that explains why generation failed, if applicable.""" progress: int | None = None @@ -2436,15 +2440,15 @@ class OpenAIVideoObject(BaseModel): model: str | None = None """The video generation model that produced the job.""" - _hidden_params: dict[str, Any] = {} + _hidden_params: dict[str, _JsonValue] = {} def __contains__(self, key) -> bool: return hasattr(self, key) - def get(self, key, default=None): + def get(self, key, default=None) -> _JsonValue: return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key) -> _JsonValue: return getattr(self, key) def json(self, **kwargs): diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index dfc45ccf9bf..88869a1edfb 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -251,8 +251,12 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" # Sample-count ceiling written on every new job: a zero-cost error loop (a shadow arm that # fails before billing) never consumes spend budget, so it must terminate on count instead. +# A multi-router job writes one attempt row per router arm, so the valve is reached +# proportionally sooner; it is a safety valve, not a sample budget. SHADOW_EVAL_TURN_VALVE: Final[int] = 10_000 +SHADOW_EVAL_MAX_ROUTERS: Final[int] = 4 + class StartShadowEvalRequest(BaseModel): """Start duplicating one or more targets' traffic for blind comparison against an auto-router. @@ -288,7 +292,24 @@ class StartShadowEvalRequest(BaseModel): "to across all their teams: JWT requests carrying their subject claim and virtual keys they own" ), ) - router_name: str = Field(description="The auto-router under evaluation, in either direction") + router_name: str | None = Field( + default=None, + description=( + "The auto-router under evaluation, in either direction: the single-router spelling of " + "router_names. Provide exactly one of the two fields" + ), + ) + router_names: tuple[str, ...] = Field( + default=(), + max_length=SHADOW_EVAL_MAX_ROUTERS, + description=( + "The auto-routers under evaluation, at most " + f"{SHADOW_EVAL_MAX_ROUTERS}. Every sampled request runs through every router listed and each " + "arm is judged independently against the same real response, so routers compare head-to-head " + "on identical traffic. More than one router requires direction 'forward'. After validation " + "this field always carries the full deduplicated set, whichever spelling the caller used" + ), + ) direction: ShadowEvalDirection = Field( default="forward", description=( @@ -332,7 +353,8 @@ class StartShadowEvalRequest(BaseModel): "Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with " "the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval " "spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight " - "samples can overshoot the cap by one sampling cache window" + "samples can overshoot the cap by one sampling cache window. Every router arm draws from the " + "same per-target budget, so a multi-router job reaches it proportionally sooner" ), ) @@ -373,6 +395,23 @@ class StartShadowEvalRequest(BaseModel): raise ValueError("baseline_model is only meaningful when direction is 'reverse'") return self + @model_validator(mode="after") + def _resolve_router_set(self) -> "StartShadowEvalRequest": + """Whichever spelling the caller used, router_names leaves validation as the full + deduplicated set, so every downstream reader consumes one field.""" + if (self.router_name is None) == (not self.router_names): + raise ValueError("provide exactly one of router_name or router_names") + single: Final = () if self.router_name is None else (self.router_name,) + routers: Final = tuple(dict.fromkeys(self.router_names or single)) + if not all(name.strip() for name in routers): + raise ValueError("router names must be non-empty strings") + if len(routers) > 1 and self.direction == "reverse": + raise ValueError("a reverse job evaluates one router against baseline_model; pass a single router") + # A returned model_copy is ignored on the __init__ construction path, so the + # normalization must land as a self attribute store to hold for every caller. + self.router_names = routers + return self + class ShadowEvalSlice(BaseModel): """Judge outcomes for one slice of a job's verdicts: a router tier, one of the @@ -428,15 +467,28 @@ class ShadowEvalResult(BaseModel): "and in reverse the models the router itself picked" ) ) + by_router: tuple[ShadowEvalSlice, ...] = Field( + default=(), + description=( + "One slice per router arm, grouped on the router name. Every arm of a multi-router job is " + "judged against the same real responses over the same sampled requests, so these slices " + "compare routers head-to-head: like-for-like win rates and spends on identical traffic. " + "Verdicts from before arm stamping existed count toward the job's own router" + ), + ) overall_shadow_win_rate_pct: float overall_tie_rate_pct: float sampled_real_spend: float = Field( default=0.0, - description="USD the real arm billed across all judged turns, cache-served turns excluded", + description=( + "USD the real arm billed across all judged turns, cache-served turns excluded. A judged turn " + "is one (request, router arm) verdict, so a multi-router job counts the real response once per " + "arm it was judged against; per-router comparisons read by_router" + ), ) sampled_shadow_spend: float = Field( default=0.0, - description="USD the shadow arm billed across the same turns, judge excluded, like for like", + description="USD the shadow arms billed across the same turns, judge excluded, like for like", ) not_sampled_count: int | None = Field( default=None, @@ -540,7 +592,13 @@ class ShadowEvalJobResponse(BaseModel): min_length=1, description="The targets whose traffic this job evaluates, and only theirs, each with its own budget", ) - router_name: str + router_names: tuple[str, ...] = Field( + min_length=1, + description=( + "Every auto-router this job runs as a shadow arm. Multi-router jobs sample one slice of " + "traffic and judge every arm against the same real responses" + ), + ) direction: ShadowEvalDirection = "forward" baseline_model: str | None = None judge_model: str @@ -562,6 +620,13 @@ class ShadowEvalJobResponse(BaseModel): last_error: str | None = Field(default=None, description="Most recent attempt error; detail endpoint only") results: ShadowEvalResult | None = Field(default=None, description="Stratified verdicts; detail endpoint only") + @computed_field + @property + def router_name(self) -> str: + """The first router, kept for callers that predate router_names; derived so the + two fields can never disagree.""" + return self.router_names[0] + @computed_field @property def status(self) -> ShadowEvalStatus: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/alice.py b/litellm/types/proxy/guardrails/guardrail_hooks/alice.py new file mode 100644 index 00000000000..73d31673dab --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/alice.py @@ -0,0 +1,21 @@ +from pydantic import Field + +from .base import GuardrailConfigModel + + +class AliceGuardrailConfigModel(GuardrailConfigModel): + api_key: str | None = Field( + default=None, + description=("The API key for Alice. If not provided, the `ALICE_API_KEY` environment variable is checked."), + ) + api_base: str | None = Field( + default=None, + description=( + "The API base URL for Alice. If not provided, the `ALICE_API_BASE` environment " + "variable is checked, then `https://api.alice.io`." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Alice" diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index 1612ea03817..7825684cfe5 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -150,6 +150,12 @@ class SCIMGroup(SCIMResource): members: list[SCIMMember] | None = None +class SCIMPlaceholderMergeResult(BaseModel): + placeholder_user_id: str + merged_into_user_id: str + team_ids: tuple[str, ...] + + # SCIM List Response Models class SCIMListResponse(BaseModel): schemas: list[str] = ["urn:ietf:params:scim:api:messages:2.0:ListResponse"] diff --git a/litellm/types/router.py b/litellm/types/router.py index ab6c807ba20..e0957383aac 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -369,7 +369,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): @model_validator(mode="before") @classmethod - def preprocess_input_data(cls, data: Any) -> Any: + def preprocess_input_data(cls, data: object) -> object: """ Pre-process input data before validation: 1. Filter out reserved Python keywords ('self', 'params', '__class__') to prevent @@ -627,6 +627,11 @@ class AlertingConfig(BaseModel): alerting_threshold: float | None = 300 +def _resolved_annotations(model_class: type[object]) -> Mapping[str, object]: + """Resolve a class's annotations, keeping each resolved annotation opaque.""" + return get_type_hints(model_class) + + class ModelGroupInfo(BaseModel): model_group: str providers: list[str] @@ -655,7 +660,7 @@ class ModelGroupInfo(BaseModel): configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None def __init__(self, **data) -> None: - for field_name, field_type in get_type_hints(self.__class__).items(): + for field_name, field_type in _resolved_annotations(self.__class__).items(): if field_type is bool and data.get(field_name) is None: data[field_name] = False super().__init__(**data) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 55a32989b1c..5783a39b30c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -193,6 +193,38 @@ class AgenticLoopParams(TypedDict, total=False): """The LLM provider name (e.g., 'bedrock', 'anthropic')""" +class OffPeakWindow(TypedDict, total=False): + """One off-peak rule: UTC time-of-day windows, optionally restricted to weekdays. + + hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them; a window may wrap past + midnight and an equal-ended window covers the whole day. weekdays is a list of days the + rule applies on, as ISO-8601 numbers (1 = Monday .. 7 = Sunday) or English day names; + omitted means every day. The weekday is read on the calendar named by the block's + weekday_timezone. + """ + + hours_utc: ReadOnly[str | Sequence[str]] + weekdays: ReadOnly[Sequence[int | str]] + + +class OffPeakPricing(TypedDict, total=False): + """Time-windowed off-peak rates for providers that discount by time of day (e.g. DeepSeek). + + hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of them for multiple daily windows, + applying on every day of the week; a window may wrap past midnight. windows adds + day-of-week-qualified rules (e.g. weekend-only whole-day off-peak), matched as a union + with hours_utc. weekday_timezone names the IANA calendar weekdays are read on, defaulting + to UTC. Any rate left unset falls back to the standard rate. + """ + + hours_utc: ReadOnly[str | Sequence[str]] + windows: ReadOnly[Sequence[OffPeakWindow]] + weekday_timezone: ReadOnly[str] + input_cost_per_token: ReadOnly[float] + output_cost_per_token: ReadOnly[float] + cache_read_input_token_cost: ReadOnly[float] + + class ModelInfoBase(ProviderSpecificModelInfo, total=False): key: Required[str] # the key in litellm.model_cost which is returned @@ -225,6 +257,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. prompt_cache_min_tokens: int | None + off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models @@ -3486,17 +3519,22 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): return {k: v for k, v in model_info.items() if k not in cls.model_fields} -SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = frozenset( - ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__ -) - frozenset(CustomPricingLiteLLMParams.model_fields) +DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) + +SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( + frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) + - frozenset(CustomPricingLiteLLMParams.model_fields) + - DEPLOYMENT_SCOPED_PRICING_FIELDS +) def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]: """Return only the fields safe to register under a shared ``{provider}/{model}`` key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus - per-deployment pricing overrides. Per-deployment metadata (``id``, - ``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key; - it stays under the deployment's unique model id. + per-deployment pricing overrides and deployment-scoped pricing blocks such as + ``off_peak_pricing``. Per-deployment metadata (``id``, ``access_via_team_ids``, + arbitrary custom keys) never belongs on the shared key; it stays under the + deployment's unique model id. """ return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS} @@ -3598,6 +3636,8 @@ all_litellm_params = ( "client", "rpm", "tpm", + "default_api_key_rpm_limit", + "default_api_key_tpm_limit", "itpm", "otpm", "max_parallel_requests", @@ -3771,6 +3811,8 @@ class LlmProviders(str, Enum): CODESTRAL = "codestral" TEXT_COMPLETION_CODESTRAL = "text-completion-codestral" DASHSCOPE = "dashscope" + QWENCLOUD = "qwencloud" + QWEN_AI_PLATFORM = "qwen_ai_platform" MODELSCOPE = "modelscope" MOONSHOT = "moonshot" PUBLICAI = "publicai" diff --git a/litellm/utils.py b/litellm/utils.py index f5adca8f272..252b6756937 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2869,10 +2869,9 @@ def _update_dictionary(existing_dict: dict, new_dict: dict) -> dict: elif isinstance(v, dict): existing_nested_dict = existing_dict.get(k) if isinstance(existing_nested_dict, dict): - existing_nested_dict.update(v) - existing_dict[k] = existing_nested_dict + existing_dict[k] = {**existing_nested_dict, **v} # mutable-ok: copy-on-write merge else: - existing_dict[k] = v + existing_dict[k] = dict(v) # mutable-ok: detached copy, never the caller's dict by reference else: existing_dict[k] = v @@ -3561,10 +3560,10 @@ def get_optional_params_embeddings( non_default_params=non_default_params, optional_params={}, kwargs=kwargs ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "gemini": - # OpenAI SDKs (and litellm's own client) send encoding_format="float" - # by default; float lists are exactly what the vertex API returns, so - # the param is a no-op — don't reject the provider default. Other - # values (e.g. "base64") stay on the unsupported-param path below. + # OpenAI SDKs send encoding_format="float" by default; float lists are + # exactly what the vertex API returns, so the param is a no-op and the + # provider default is not rejected. Other values (e.g. "base64") stay + # on the unsupported-param path below. if non_default_params.get("encoding_format") == "float": non_default_params.pop("encoding_format") supported_params = get_supported_openai_params( @@ -4877,11 +4876,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None: def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: if target_order is not None: - filtered: Final = [d for d in healthy_deployments if _get_deployment_order(d) == target_order] - if filtered: - return filtered - # target_order doesn't match any deployment (e.g., external fallback model) — return all - return healthy_deployments + return [d for d in healthy_deployments if _get_deployment_order(d) == target_order] # Default: pick min order group _valid_orders: Final[list[int]] = [ @@ -5860,6 +5855,7 @@ def _get_model_info_helper( cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), + off_peak_pricing=_model_info.get("off_peak_pricing", None), input_cost_per_character=_model_info.get("input_cost_per_character", None), input_cost_per_token_above_128k_tokens=_model_info.get("input_cost_per_token_above_128k_tokens", None), input_cost_per_token_above_200k_tokens=_model_info.get("input_cost_per_token_above_200k_tokens", None), @@ -6586,11 +6582,11 @@ def validate_environment( keys_in_environment = True else: missing_keys.append("WANDB_API_KEY") - elif custom_llm_provider == "dashscope": - if "DASHSCOPE_API_KEY" in os.environ: + elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"): + if f"{custom_llm_provider.upper()}_API_KEY" in os.environ or "DASHSCOPE_API_KEY" in os.environ: keys_in_environment = True else: - missing_keys.append("DASHSCOPE_API_KEY") + missing_keys.append(f"{custom_llm_provider.upper()}_API_KEY") elif custom_llm_provider == "modelscope": if "MODELSCOPE_API_KEY" in os.environ: keys_in_environment = True @@ -8152,6 +8148,11 @@ class ProviderConfigManager: LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False), LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False), LlmProviders.DASHSCOPE: (lambda: litellm.DashScopeChatConfig(), False), + LlmProviders.QWENCLOUD: (lambda: litellm.QwenCloudChatConfig(), False), + LlmProviders.QWEN_AI_PLATFORM: ( + lambda: litellm.QwenAIPlatformChatConfig(), + False, + ), LlmProviders.MODELSCOPE: (lambda: litellm.ModelScopeChatConfig(), False), LlmProviders.MOONSHOT: (lambda: litellm.MoonshotChatConfig(), False), LlmProviders.DOCKER_MODEL_RUNNER: ( @@ -8278,10 +8279,17 @@ class ProviderConfigManager: """ # Handle OpenAI special cases (O-series and GPT-5 models) if provider == LlmProviders.OPENAI: + from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIGPTConfig, + OpenAIUnknownModelConfig, + ) + if litellm.openaiOSeriesConfig.is_model_o_series_model(model=model): return litellm.openaiOSeriesConfig if litellm.OpenAIGPT5Config.is_model_gpt_5_model(model=model): return litellm.OpenAIGPT5Config() + if not OpenAIGPTConfig.is_openai_catalog_model(model): + return OpenAIUnknownModelConfig() # Handle Azure before the generic map so base_model can be threaded through if provider == LlmProviders.AZURE: @@ -8359,12 +8367,16 @@ class ProviderConfigManager: ) return VolcEngineEmbeddingConfig() - elif litellm.LlmProviders.DASHSCOPE == provider: - from litellm.llms.dashscope.embed.transformation import ( - DashScopeEmbeddingConfig, + elif provider in ( + litellm.LlmProviders.DASHSCOPE, + litellm.LlmProviders.QWENCLOUD, + litellm.LlmProviders.QWEN_AI_PLATFORM, + ): + from litellm.llms.dashscope.common_utils import ( + get_dashscope_family_embedding_config, ) - return DashScopeEmbeddingConfig() + return get_dashscope_family_embedding_config(provider.value) elif litellm.LlmProviders.OVHCLOUD == provider: return litellm.OVHCloudEmbeddingConfig() elif litellm.LlmProviders.SNOWFLAKE == provider: @@ -8437,12 +8449,16 @@ class ProviderConfigManager: return litellm.VoyageRerankConfig() elif litellm.LlmProviders.WATSONX == provider: return litellm.IBMWatsonXRerankConfig() - elif litellm.LlmProviders.DASHSCOPE == provider: - from litellm.llms.dashscope.rerank.transformation import ( - DashScopeRerankConfig, + elif provider in ( + litellm.LlmProviders.DASHSCOPE, + litellm.LlmProviders.QWENCLOUD, + litellm.LlmProviders.QWEN_AI_PLATFORM, + ): + from litellm.llms.dashscope.common_utils import ( + get_dashscope_family_rerank_config, ) - return DashScopeRerankConfig() + return get_dashscope_family_rerank_config(provider.value) return litellm.CohereRerankConfig() @staticmethod @@ -9115,12 +9131,16 @@ class ProviderConfigManager: ) return get_openrouter_image_generation_config(model) - elif LlmProviders.DASHSCOPE == provider: - from litellm.llms.dashscope.image_generation import ( - get_dashscope_image_generation_config, + elif provider in ( + LlmProviders.DASHSCOPE, + LlmProviders.QWENCLOUD, + LlmProviders.QWEN_AI_PLATFORM, + ): + from litellm.llms.dashscope.common_utils import ( + get_dashscope_family_image_generation_config, ) - return get_dashscope_image_generation_config(model) + return get_dashscope_family_image_generation_config(provider.value) elif LlmProviders.MODELSCOPE == provider: from litellm.llms.modelscope.image_generation import ( get_modelscope_image_generation_config, @@ -9433,6 +9453,10 @@ class ProviderConfigManager: return RunwayMLTextToSpeechConfig() elif litellm.LlmProviders.VERTEX_AI == provider: + if "gemini" in model: + # Gemini TTS uses the speech_to_completion bridge, and Google Cloud TTS param + # mapping would drop response_format before the bridge sees it (LIT-6501) + return None from litellm.llms.vertex_ai.text_to_speech.transformation import ( VertexAITextToSpeechConfig, ) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 22d27bc3266..b71d6784873 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -112,7 +112,9 @@ class VectorStoreRegistry: Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS. """ # Get the list of supported param names from the Literal type - supported_params: Final = get_args(VECTOR_STORE_OPENAI_PARAMS) + supported_params: Final = tuple( + param for param in get_args(VECTOR_STORE_OPENAI_PARAMS) if isinstance(param, str) + ) # Extract only the params that exist in the tool kwargs: Final = {param: tool.get(param) for param in supported_params if param in tool} @@ -503,7 +505,7 @@ class VectorStoreRegistry: vector_stores_from_db.append(_litellm_managed_vector_store) return vector_stores_from_db - def get_credentials_for_vector_store(self, vector_store_id: str) -> dict[str, Any]: + def get_credentials_for_vector_store(self, vector_store_id: str) -> dict[str, object]: """ Get the credentials for a vector store diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 718e6c489fd..10e8b56d925 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -554,6 +554,7 @@ "supports_vision": true }, "amazon.nova-sonic-v1:0": { + "deprecation_date": "2026-09-14", "input_cost_per_audio_token": 3.4e-06, "input_cost_per_token": 6e-08, "litellm_provider": "bedrock", @@ -1450,6 +1451,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, @@ -1487,6 +1526,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "global.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -1524,6 +1601,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "us.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 2.75e-07, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "eu.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -1561,6 +1676,44 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "eu.anthropic.claude-fable-5-1": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 2.75e-07, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -1591,7 +1744,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1627,7 +1780,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1663,7 +1816,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1699,7 +1852,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1735,7 +1888,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -1771,7 +1924,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "supports_parallel_tool_use_config": true, @@ -2064,7 +2217,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2101,7 +2254,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2138,7 +2291,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2175,7 +2328,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2212,7 +2365,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2249,7 +2402,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_native_structured_output": true, + "supports_native_structured_output": false, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -3045,6 +3198,7 @@ "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { + "deprecation_date": "2027-12-05", "supports_mid_conversation_system": true, "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -3077,7 +3231,43 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "azure_ai/claude-fable-5-1": { + "supports_mid_conversation_system": true, + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "azure_ai/claude-opus-5": { + "deprecation_date": "2027-07-08", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3110,6 +3300,7 @@ "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-4-8": { + "deprecation_date": "2027-09-01", "supports_mid_conversation_system": true, "supports_adaptive_thinking": true, "input_cost_per_token": 5e-06, @@ -3188,6 +3379,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-5": { + "deprecation_date": "2027-06-30", "supports_mid_conversation_system": true, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, @@ -9451,7 +9643,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Flash": { "input_cost_per_image_token": 1.75e-06, @@ -9464,7 +9657,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", @@ -9963,7 +10157,9 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.45e-07, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", @@ -9977,6 +10173,24 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.8e-08, + "supports_prompt_caching": true + }, + "azure_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 1.4e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 4.4e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, "supports_tool_choice": true }, "azure_ai/embed-v-4-0": { @@ -10192,11 +10406,13 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supports_function_calling": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1e-07, + "supports_prompt_caching": true }, "azure_ai/kimi-k2.6": { "deprecation_date": "2027-04-16", @@ -10207,7 +10423,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k2-6-in-microsoft-foundry/4513125", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supported_modalities": [ "text", "image" @@ -10218,7 +10434,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.6e-07, + "supports_prompt_caching": true }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, @@ -11902,7 +12120,7 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 2.65e-06, + "output_cost_per_token": 6e-07, "supports_pdf_input": true }, "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { @@ -12287,6 +12505,7 @@ "supports_tool_choice": true }, "cerebras/zai-glm-4.7": { + "deprecation_date": "2026-08-17", "input_cost_per_token": 2.25e-06, "litellm_provider": "cerebras", "max_input_tokens": 128000, @@ -13022,6 +13241,47 @@ "supports_native_structured_output": true, "source": "https://docs.anthropic.com/en/docs/about-claude/models/overview" }, + "claude-fable-5-1": { + "deprecation_date": "2027-09-01", + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/models/fable-5-1/overview" + }, "claude-opus-5": { "deprecation_date": "2027-07-24", "cache_creation_input_token_cost": 6.25e-06, @@ -14718,6 +14978,1910 @@ "/v1/images/generations" ] }, + "qwencloud/deepseek-v4-flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/deepseek-v4-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/glm-5.1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 202745, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/glm-5.2": { + "cache_read_input_token_cost": 2.8e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 229376, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen-coder": { + "input_cost_per_token": 3e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-flash": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-flash-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-max": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 30720, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 6.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-01-25": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-04-28": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-07-14": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-plus-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-plus-2025-09-11": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-plus-latest": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen-turbo": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-2024-11-01": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-2025-04-28": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-turbo-latest": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwencloud", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-30b-a3b": { + "litellm_provider": "qwencloud", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-coder-flash": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-flash-2025-07-28": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-coder-plus-2025-07-22": { + "litellm_provider": "qwencloud", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3-max-preview": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-max": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-max-2026-01-23": { + "litellm_provider": "qwencloud", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwencloud/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwen3-vl-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "qwencloud/qwen3.5-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3.7-max": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwencloud/qwen3.7-plus": { + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwencloud/qwen3.8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "qwencloud", + "max_input_tokens": 991808, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwencloud/qwq-plus": { + "input_cost_per_token": 8e-07, + "litellm_provider": "qwencloud", + "max_input_tokens": 98304, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://www.qwencloud.com/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwencloud/qwen-image-2.0": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-2.0-pro": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-3.0": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwencloud/qwen-image-3.0-pro": { + "litellm_provider": "qwencloud", + "mode": "image_generation", + "source": "https://www.qwencloud.com/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/deepseek-v4-flash": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/deepseek-v4-pro": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 4.8e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/glm-5.1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 202745, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/glm-5.2": { + "cache_read_input_token_cost": 2.8e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 229376, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen-coder": { + "input_cost_per_token": 3e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-flash": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-flash-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 4e-07, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 2e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-max": { + "input_cost_per_token": 1.6e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 30720, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 6.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-01-25": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-04-28": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-07-14": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-plus-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-plus-2025-09-11": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-plus-latest": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 4e-06, + "output_cost_per_token": 1.2e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 3.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen-turbo": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-2024-11-01": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-2025-04-28": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-turbo-latest": { + "input_cost_per_token": 5e-08, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 1000000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_reasoning_token": 5e-07, + "output_cost_per_token": 2e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-30b-a3b": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 129024, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-coder-flash": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-flash-2025-07-28": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 8e-07, + "output_cost_per_token": 4e-06, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 1.6e-06, + "output_cost_per_token": 9.6e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "cache_read_input_token_cost": 1.8e-07, + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-coder-plus-2025-07-22": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 997952, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 1.8e-06, + "output_cost_per_token": 9e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 256000.0 + ] + }, + { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 6e-05, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max-preview": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-max-2026-01-23": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 258048, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "tiered_pricing": [ + { + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "range": [ + 128000.0, + 252000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3-next-80b-a3b-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-next-80b-a3b-thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3-vl-235b-a22b-instruct": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-235b-a22b-thinking": { + "input_cost_per_token": 4e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-32b-instruct": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.4e-07, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-32b-thinking": { + "input_cost_per_token": 1.6e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.87e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwen3-vl-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 260096, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 2e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 32000.0 + ] + }, + { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 32000.0, + 128000.0 + ] + }, + { + "input_cost_per_token": 6e-07, + "output_cost_per_token": 4.8e-06, + "range": [ + 128000.0, + 256000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.5-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2.4e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 3e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.7-max": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen3.7-plus": { + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tiered_pricing": [ + { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, + "range": [ + 0, + 256000.0 + ] + }, + { + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 4.8e-06, + "range": [ + 256000.0, + 1000000.0 + ] + } + ] + }, + "qwen_ai_platform/qwen3.8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 991808, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "qwen_ai_platform/qwq-plus": { + "input_cost_per_token": 8e-07, + "litellm_provider": "qwen_ai_platform", + "max_input_tokens": 98304, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "qwen_ai_platform/qwen-image-2.0": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-2.0-pro": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-3.0": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "qwen_ai_platform/qwen-image-3.0-pro": { + "litellm_provider": "qwen_ai_platform", + "mode": "image_generation", + "source": "https://www.alibabacloud.com/help/en/model-studio/models", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "databricks/databricks-bge-large-en": { "cache_creation_input_token_cost": 1.0003e-07, "cache_read_input_token_cost": 1.0003e-07, @@ -15101,6 +17265,62 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/databricks-deepseek-v4-flash-0731": { + "cache_creation_input_token_cost": 1.4e-07, + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "input_dbu_cost_per_token": 2e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "output_dbu_cost_per_token": 4e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "databricks/databricks-deepseek-v4-pro-0813": { + "cache_creation_input_token_cost": 1.31999e-06, + "cache_read_input_token_cost": 1.3202e-07, + "input_cost_per_token": 1.31999e-06, + "input_dbu_cost_per_token": 1.8857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 393216, + "max_tokens": 393216, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference. Context/max output are the DeepSeek-published model limits (1M context, 384K max output)." + }, + "mode": "chat", + "output_cost_per_token": 3.95997e-06, + "output_dbu_cost_per_token": 5.6571e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": false + }, "databricks/databricks-gemini-2-5-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.0002e-08, @@ -20631,6 +22851,7 @@ "supports_image_size": false }, "gemini-live-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", @@ -21303,6 +23524,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -23140,6 +25418,65 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -23548,6 +25885,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, @@ -23852,8 +26246,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23867,7 +26263,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23895,8 +26292,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -23910,7 +26309,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", + "output_cost_per_second_4k": 0.6, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -26881,16 +29281,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -26902,6 +29305,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -26944,16 +29348,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -26965,6 +29372,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -27008,16 +29416,19 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -27029,6 +29440,7 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -27071,16 +29483,19 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -27092,6 +29507,7 @@ "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, @@ -27331,7 +29747,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -27385,7 +29804,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -27534,7 +29956,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -27583,7 +30008,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -27632,7 +30060,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, @@ -27681,7 +30111,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -28617,17 +31049,18 @@ }, "gpt-realtime-2": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, "litellm_provider": "openai", - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, - "output_cost_per_token": 1.6e-05, + "output_cost_per_token": 2.4e-05, "supported_endpoints": [ "/v1/realtime" ], @@ -28691,8 +31124,8 @@ "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, @@ -28724,7 +31157,7 @@ "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "realtime", @@ -31488,19 +33921,21 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/magistral-medium-latest": { - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 5e-06, - "source": "https://mistral.ai/news/magistral", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-2506": { "deprecation_date": "2025-11-30", @@ -31519,19 +33954,21 @@ "supports_tool_choice": true }, "mistral/magistral-small-latest": { - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://mistral.ai/pricing#api-pricing", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-1-2-2509": { "deprecation_date": "2026-07-31", @@ -31663,16 +34100,21 @@ "supports_vision": true }, "mistral/mistral-medium": { - "input_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.1e-06, + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-medium-2312": { "deprecation_date": "2025-06-16", @@ -36339,12 +38781,22 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, "parallel_ai/search": { - "input_cost_per_query": 0.004, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-fast": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, "parallel_ai/search-pro": { - "input_cost_per_query": 0.009, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-turbo": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, @@ -39113,13 +41565,13 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6.25e-06, + "output_cost_per_token": 6e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -39729,6 +42181,70 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "us-gov.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us-gov.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_1hr": 2.2e-06, @@ -41892,6 +44408,42 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "vertex_ai/claude-fable-5-1": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", "regional_endpoint_uplift_multiplier": 1.1, @@ -41927,6 +44479,42 @@ "supports_max_reasoning_effort": true, "prompt_cache_min_tokens": 512 }, + "vertex_ai/claude-fable-5-1@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512 + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, @@ -43440,7 +46028,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43453,8 +46042,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43469,7 +46060,8 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second_4k": 0.6, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43483,8 +46075,10 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.12, + "output_cost_per_second_4k": 0.3, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_modalities": [ "text" ], @@ -43544,6 +46138,26 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "voyage/rerank-3": { + "input_cost_per_token": 5e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/rerank-3-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-2": { "input_cost_per_token": 1e-07, "litellm_provider": "voyage", @@ -44795,6 +47409,27 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-build-latest": { + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/developers/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, @@ -52231,14 +54866,14 @@ "supports_vision": true }, "fireworks_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "input_cost_per_token": 1.4e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.8e-07, + "output_cost_per_token": 6.6e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -55115,6 +57750,34 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, @@ -55124,5 +57787,592 @@ "max_tokens": 40960, "mode": "embedding", "source": "https://docs.fireworks.ai/serverless/pricing" + }, + "zai/glm-5.2": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.z.ai/guides/overview/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3.8-Flash": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 4.7e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "cerebras/gemma-4-31b": { + "input_cost_per_token": 9.9e-07, + "litellm_provider": "cerebras", + "max_input_tokens": 131072, + "max_output_tokens": 40960, + "max_tokens": 40960, + "mode": "chat", + "output_cost_per_token": 1.49e-06, + "source": "https://api.cerebras.ai/public/v1/models/gemma-4-31b", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "elevenlabs/scribe_v2": { + "input_cost_per_second": 6.11e-05, + "litellm_provider": "elevenlabs", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://elevenlabs.io/pricing/api", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "scaleway/glm-5.2": { + "input_cost_per_token": 1.8e-06, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false + }, + "scaleway/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": false + }, + "azure_ai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "deprecation_date": "2026-10-03", + "input_cost_per_token": 9.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-west-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { + "use_openai_responses_path": true, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://aws.amazon.com/bedrock/pricing/", + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2.4e-07 + }, + "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "azure/us-gov/gpt-5.1": { + "cache_read_input_token_cost": 1.71875e-07, + "default_reasoning_effort": "none", + "input_cost_per_token": 1.71875e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.375e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure/us-gov/o3-mini": { + "cache_read_input_token_cost": 7.57e-07, + "input_cost_per_token": 1.513e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.05e-06, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "azure/us-gov/text-embedding-3-large": { + "input_cost_per_token": 1.63e-07, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "azure/us-gov/text-embedding-3-small": { + "input_cost_per_token": 2.5e-08, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/openai/whisper": { + "input_cost_per_second": 7.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "cloudflare/@cf/openai/whisper-large-v3-turbo": { + "input_cost_per_second": 8.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 3f6d3b4f910..9e370e5406a 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -662,6 +662,9 @@ "supports_embedding_image_input": { "type": "boolean" }, + "supports_forced_tool_use": { + "type": "boolean" + }, "supports_function_calling": { "type": "boolean" }, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 7c7d508856f..ebc220b3496 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -724,6 +724,42 @@ "interactions": true } }, + "qwencloud": { + "display_name": "QwenCloud (`qwencloud`)", + "url": "https://docs.litellm.ai/docs/providers/qwencloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "qwen_ai_platform": { + "display_name": "Qwen AI Platform (`qwen_ai_platform`)", + "url": "https://docs.litellm.ai/docs/providers/qwencloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, "databricks": { "display_name": "Databricks (`databricks`)", "url": "https://docs.litellm.ai/docs/providers/databricks", diff --git a/pyproject.toml b/pyproject.toml index f33bf68822a..2866e27e84c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.100.0" +version = "1.101.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -67,8 +67,8 @@ 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.91", - "litellm-enterprise==0.1.62", + "litellm-proxy-extras==0.4.92", + "litellm-enterprise==0.1.63", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -319,7 +319,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.100.0" +version = "1.101.0" version_files = [ "pyproject.toml:^version", ] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 1e855ea1b71..9b1cc977a64 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2995 + "limit": 2985 }, "ANN002": { "limit": 71 @@ -9,13 +9,13 @@ "limit": 809 }, "ANN201": { - "limit": 2002 + "limit": 2001 }, "ANN202": { - "limit": 845 + "limit": 835 }, "ANN204": { - "limit": 698 + "limit": 693 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 587 + "limit": 307 }, "ASYNC230": { "limit": 11 @@ -117,7 +117,7 @@ "limit": 1 }, "PERF102": { - "limit": 23 + "limit": 21 }, "PERF401": { "limit": 12 @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 175 + "limit": 173 }, "RUF012": { "limit": 239 @@ -198,7 +198,7 @@ "limit": 56 }, "SIM102": { - "limit": 314 + "limit": 310 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1108 + "limit": 1073 }, "TRY002": { "limit": 524 @@ -240,13 +240,13 @@ "limit": 96 }, "TRY201": { - "limit": 405 + "limit": 403 }, "TRY203": { - "limit": 113 + "limit": 111 }, "TRY300": { - "limit": 855 + "limit": 854 }, "UP028": { "limit": 2 diff --git a/ruff-strict.toml b/ruff-strict.toml index 7afc5da71ee..ae092bdde7d 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -26,6 +26,10 @@ external = [ # caught a real mismatch, confirming Any is correct here, not a shortcut. "litellm/litellm_core_utils/litellm_logging.py" = ["ANN401"] "litellm/utils.py" = ["ANN401"] +# `**kwargs` forwards verbatim to CustomGuardrail.__init__, whose param list is wide and +# grows over time; typing it concretely (`object`) broke that forwarding call outright — +# basedpyright turned every named param into a reportArgumentType error. Any is correct here. +"litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"] [lint.mccabe] max-complexity = 15 diff --git a/schema.prisma b/schema.prisma index a9c468c3e9d..b44935ce3dd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1533,7 +1533,8 @@ model LiteLLM_ShadowEvalJob { group_id String // legs of one job share this; the API's job id target_type String @default("key") // key | team | user target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows - router_name String // the auto-router under evaluation, in either direction + router_name String // first (often only) auto-router under evaluation; router_names is the full set + router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name) direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String @@ -1557,6 +1558,7 @@ model LiteLLM_ShadowEvalAttempt { job_id String request_id String // the judged real request outcome String // real | shadow | tie | error + router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router tier String? // router's tier for the prompt, when classified real_model String? shadow_model String? diff --git a/test-quality-budget.json b/test-quality-budget.json index ee33eb581d6..d834c581609 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 742 + "limit": 741 }, "TQ003": { "limit": 62 @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11139 + "limit": 11135 } } diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 37e940460f6..790956156b0 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -63,6 +63,7 @@ IGNORE_FUNCTIONS = [ "json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks. "_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). + "_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input. ] diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index c541c035db7..a5e00799519 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -81,6 +81,7 @@ ignored_function_names = [ "_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name) "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call + "_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name) ] diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 1e6de0c3d6a..d571fb36546 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -64,14 +64,12 @@ - {id: mgmt.budget.list_v1.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "management_v1/budgets.py:129", rationale: "Budget enumeration the Budgets page can page, sort and filter"} - {id: mgmt.budget.list_v1.admin_only, module: mgmt, tier: P1, surface: api, assertions: [admin_only], source: "management_v1/budgets.py:129", rationale: "A caller without admin view is refused, not served an empty page"} - {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"} -- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."} - {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"} - {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"} - {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"} - {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"} - {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"} - {id: mgmt.fallback_management.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "fallback_management_endpoints.py", rationale: "Fallback config (smoke)"} -- {id: mgmt.config_override.hashicorp_vault.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "config_override_endpoints.py", rationale: "Vault integration (smoke)"} - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index e4f59f5c4d2..c5971c5362c 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -2,10 +2,16 @@ The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report (`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records -outcome, duration, and node id for every ``; the only signals it cannot -derive on its own are the normalized suite package and the coverage-registry cell -ids a test covers. Those ride along as JUnit `` entries via each item's -`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`. +outcome, duration, and node id for every ``; the signals it cannot +derive on its own are the normalized suite package, the coverage-registry cell +ids a test covers, and where the test's source lives. Those ride along as JUnit +`` entries via each item's `user_properties`, attached in +`conftest.py::pytest_collection_modifyitems`. + +`source` is a property rather than the `file=` / `line=` attributes pytest used +to write, because the `xunit2` family this suite runs on drops those, and +switching families would change the XML for every consumer of it -- the +Buildkite Test Engine upload and the Loki pipeline included. """ from __future__ import annotations @@ -14,22 +20,61 @@ from collections.abc import Iterable import pytest +# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing +# at runtime names this suite's place in the repo. test_junit_properties.py +# fails from a checkout if it moves. +SUITE_ROOT = "tests/e2e" + + +def suite_parts(path_part: str) -> tuple[str, ...]: + """Path components of a suite file relative to tests/e2e, however it ran. + + Pytest paths are rootdir-relative, and rootdir moves with the invocation: a + repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the + runner image) gives `logging/test_x.py`. Both collapse to the same tuple. + """ + raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".") + return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + def package_from_nodeid(nodeid: str) -> str: - """Top-level suite package under tests/e2e/, or 'root' for top-level files. - - Pytest nodeids are relative to the invocation cwd. Repo-root runs look like - `tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the - `tests/e2e` prefix so package is the suite dir either way. - """ - path_part = nodeid.split("::", 1)[0].replace("\\", "/") - raw = tuple(p for p in path_part.split("/") if p and p != ".") - parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + """Top-level suite package under tests/e2e/, or 'root' for top-level files.""" + parts = suite_parts(nodeid.split("::", 1)[0]) if len(parts) <= 1: return "root" return parts[0] +def source_from_location(path: str, lineno: int | None) -> str: + """Repo-relative `path:line` for a test, or '' when nothing is linkable. + + `pytest.Item.location` gives a rootdir-relative path and a ZERO-based line. + The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was + started, and the line is emitted ONE-based to match editors, tracebacks and + code hosts. A decorated test anchors at its first decorator, which is where + pytest reports it. + + Empty rather than a guess for anything unlinkable: no line, a path reaching + upward, or a path carrying a colon, which is both how an absolute Windows + path arrives and a character `path:line` has no way to represent. + """ + if lineno is None: + return "" + normalized = path.replace("\\", "/") + if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"): + return "" + parts = suite_parts(normalized) + if not parts: + return "" + return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}" + + +def source_from_item(item: pytest.Item) -> str: + """Read the repo-relative `path:line` off a pytest Item's reported location.""" + path, lineno, _ = item.location + return source_from_location(path, lineno) + + def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]: """Flatten @pytest.mark.covers arg lists into unique, order-preserving cell ids, dropping anything that is not a non-empty string.""" @@ -43,10 +88,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]: def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: """The custom signals a standard reporter cannot derive: the normalized suite - package and the comma-joined coverage-registry cell ids this test covers.""" + package, the comma-joined coverage-registry cell ids this test covers, and the + repo-relative `path:line` its source sits at.""" return ( ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), + ("source", source_from_item(item)), ) diff --git a/tests/e2e/management/test_config_misc_endpoints_e2e.py b/tests/e2e/management/test_config_misc_endpoints_e2e.py index 195732c0201..099ffa4b3bd 100644 --- a/tests/e2e/management/test_config_misc_endpoints_e2e.py +++ b/tests/e2e/management/test_config_misc_endpoints_e2e.py @@ -7,9 +7,13 @@ so a read-back reflects the change. Router settings, which mutate global proxy state, are exercised with a benign, self-restoring change so a shared proxy is left as it was found. -Cache settings are deliberately not covered here; see the rationale on -mgmt.cache_settings.update.happy_path in coverage_registry/mgmt.yaml before adding -a test for that route. +Cache settings and the Vault config override are deliberately not covered here. +Both routes reconfigure the whole proxy: /cache/settings persists what it receives +into a row that outranks the YAML cache_params and is re-applied on a timer, and +/config_overrides/hashicorp_vault swaps the process-wide secret manager. Neither can +be exercised safely against the shared proxy the suites run on, so they need an +isolated proxy before a test lands. Do not add a read-then-write-back test for +either one. """ from __future__ import annotations diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 88ab5666084..68005ae3f6a 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -30,6 +30,34 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str return key +class TestMcpKeyGrantByAlias: + def test_alias_grant_persists_verbatim_and_lists_tools( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + """A key granted an MCP server by its alias must store the alias, not the + resolved server_id: in a shared-DB multi-region deployment each instance + derives a different id for the same config server, so only the alias + grants access on every region. The same key must still see the server's + tools, proving the alias grant is honored at request time.""" + server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) + alias = next(row.alias for row in client.registered_servers() if row.server_id == server_id) + assert alias, f"registered server {server_id} has no alias to grant by" + + key = _key(client, resources, mcp_servers=[alias]) + + stored = client.proxy.key_info(key).object_permission + assert stored is not None and stored.mcp_servers == [alias], ( + f"alias grant was rewritten before persisting (expected [{alias!r}]): " + f"{stored.mcp_servers if stored else None}. A stored server_id is region-local " + f"and breaks the grant on every other instance sharing this database" + ) + + _ = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) + + class TestMcpKeyWithoutAccessIsDenied: @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission") def test_list_tools_denied_without_permission( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 6d9ccad9a24..79d9e011f7e 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -114,6 +114,7 @@ class KeyInfo(BaseModel): budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None budget_limits: list[BudgetWindowState] | None = None + object_permission: ObjectPermission | None = None class KeyInfoResponse(BaseModel): @@ -805,6 +806,7 @@ class LiteLLMParamsBody(BaseModel): mock_response: str | None = None timeout: float | None = None tpm: int | None = None + weight: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] @@ -819,6 +821,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails_policy: dict[str, int] | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index e342aa363ca..5822058003c 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -19,6 +19,8 @@ from models import ( ChatMessage, ChatResponse, LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, ReliabilityChatBody, RouterSettingsOverride, ) @@ -26,6 +28,18 @@ from models import ( REAL_MODEL = "openai/gpt-5.5" REAL_KEY = "os.environ/OPENAI_API_KEY" +# The smallest-context chat model OpenAI still serves (16385 tokens). A prompt +# past that limit comes back as a real `context_length_exceeded` 400, which is +# what litellm maps to ContextWindowExceededError. +SMALL_CONTEXT_MODEL = "openai/gpt-3.5-turbo" +SMALL_CONTEXT_LIMIT_TOKENS = 16385 + + +def oversized_prompt(marker: str) -> str: + """A prompt comfortably past SMALL_CONTEXT_MODEL's context limit, so the + provider refuses it on length rather than answering a truncated version.""" + return f"{marker} " + ("token " * (SMALL_CONTEXT_LIMIT_TOKENS + 4000)) + def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: """Register a deployment pointing at an unreachable base, so every call to it @@ -40,6 +54,38 @@ def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) +def create_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment on the smallest-context model OpenAI still serves, so an + oversized prompt earns a real context-window refusal from the provider.""" + return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY)) + + +def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair: a 1ms deadline the backend always + exceeds, all of the model group's shuffle weight, and a cooldown policy that + benches it on its first Timeout so the retry cannot land on it again.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001, weight=1), + model_info=ModelInfoBody(allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + ) + ) + + +def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: + """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle + never opens on it. It is reachable only once its sibling is benched and the + weighted pick falls through to a uniform one over what is left.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, weight=0), + model_info=ModelInfoBody(), + ) + ) + + def chat_override( proxy: ProxyClient, key: str, diff --git a/tests/e2e/router/test_reliability_fallbacks_e2e.py b/tests/e2e/router/test_reliability_fallbacks_e2e.py index d3ce62f8f95..8cece41ce2d 100644 --- a/tests/e2e/router/test_reliability_fallbacks_e2e.py +++ b/tests/e2e/router/test_reliability_fallbacks_e2e.py @@ -9,6 +9,10 @@ in the x-litellm-attempted-fallbacks header. Empty content is accepted only when `finish_reason == "length"` and the response billed completion tokens, since gpt-5.5 counts reasoning against max_tokens and can consume the whole budget before emitting any text; a fallback that produced nothing at all still fails. + +The context-window case is a different reroute from a plain failure: the provider +refuses the prompt on length, and `context_window_fallbacks` is the setting that +reroutes it, not `fallbacks`. """ from __future__ import annotations @@ -25,8 +29,10 @@ from reliability_support import ( completion_tokens_of, content_of, create_bad_base_deployment, + create_small_context_deployment, create_timeout_deployment, finish_reason_of, + oversized_prompt, reasoning_tokens_of, ) @@ -82,3 +88,17 @@ class TestReliabilityFallbacks: override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), ) _assert_served_by_fallback(resp) + + @pytest.mark.covers("reliability.fallback.context_window.routes_to_fallback") + def test_context_window_routes_to_fallback( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + primary = f"reliability-ctxfail-{unique_marker()}" + model_id = create_small_context_deployment(client.proxy, primary) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override( + client.proxy, scoped_key, primary, oversized_prompt(unique_marker()), + override=RouterSettingsOverride(context_window_fallbacks=[{primary: ["gpt-5.5"]}]), + ) + _assert_served_by_fallback(resp) diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py new file mode 100644 index 00000000000..5441412935c --- /dev/null +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -0,0 +1,73 @@ +"""Live e2e: a request that fails on its first deployment is retried inside its own +model group and still comes back a completion. + +The model group is a pair: an always-timing-out deployment that holds all of the +group's shuffle weight, and a healthy backup at weight 0. The weighted pick always +opens on the timing-out one, its first Timeout benches it (an +`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls +through to the only deployment left. So the customer sees a completion and the +proxy reports that it took a retry to get there, with no random first pick in the +middle of it. +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import RouterSettingsOverride +from reliability_support import ( + chat_override, + completion_tokens_of, + content_of, + create_always_timing_out_deployment, + create_zero_weight_backup_deployment, + finish_reason_of, +) + +pytestmark = pytest.mark.e2e + + +class TestReliabilityRetries: + @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") + def test_timeout_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + timing_out = create_always_timing_out_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(timing_out)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=2), + ) + + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the timing-out deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py new file mode 100644 index 00000000000..02c1413c840 --- /dev/null +++ b/tests/e2e/test_junit_properties.py @@ -0,0 +1,131 @@ +"""Harness coverage for the custom JUnit properties. + +No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree +about where a suite file lives -- ``package_from_nodeid`` (strip the suite root) +and ``source_from_location`` (re-root at it) -- across both ways the suite is +launched, plus the one-based line offset and the refusal to emit a path that +escapes the suite. The consumers of these properties are the Loki/Grafana +rollups and, for ``source``, the status page's per-test links to GitHub. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from junit_properties import ( + SUITE_ROOT, + attach_result_properties, + dedupe_covers, + package_from_nodeid, + result_properties, + source_from_location, + suite_parts, +) + + +def collected_item(request: pytest.FixtureRequest, name: str) -> pytest.Item: + """The Item pytest collected for test ``name`` in this file: the real nodeid, + location and marker machinery the collection hook reads, as pytest built it.""" + return next(item for item in request.session.items if item.path == request.path and item.name == name) + + +def repo_root() -> Path | None: + """The litellm checkout above this file, or None when there isn't one.""" + return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None) + + +class TestSuiteParts: + @pytest.mark.parametrize( + "path", + ["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"], + ) + def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None: + """A repo-root run and a suite-cwd run report the same file differently; + every downstream signal has to see one spelling.""" + assert suite_parts(path) == ("logging", "test_x.py") + + def test_top_level_suite_file_keeps_its_single_component(self) -> None: + assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",) + + +class TestPackageFromNodeid: + @pytest.mark.parametrize( + ("nodeid", "expected"), + [ + ("logging/test_x.py::TestFoo::test_bar", "logging"), + ("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"), + ("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"), + ("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"), + ("tests/e2e/test_fixture_mode.py::test_bar", "root"), + ], + ) + def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None: + assert package_from_nodeid(nodeid) == expected + + +class TestSourceFromLocation: + @pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"]) + def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None: + assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41" + + def test_line_is_emitted_one_based(self) -> None: + """pytest.Item.location counts from 0; editors, tracebacks and GitHub's + #L anchor all count from 1, and an off-by-one lands on the decorator.""" + assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1" + + def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None: + assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40" + + @pytest.mark.parametrize( + ("path", "lineno"), + [ + ("a2a/test_x.py", None), + ("/app/e2e/a2a/test_x.py", 40), + ("C:\\app\\e2e\\a2a\\test_x.py", 40), + ("../conftest.py", 40), + ("", 40), + ], + ) + def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None: + """A colon is rejected on two counts: it is how a Windows absolute path + arrives, and `path:line` cannot represent one in the path half.""" + assert source_from_location(path, lineno) == "" + + +class TestResultProperties: + def test_every_test_carries_package_covers_and_source(self, request: pytest.FixtureRequest) -> None: + """Read off this test's own collected Item, so the nodeid and location are + whatever pytest reports for the launch shape in use, and the marker is added + at run time so the coverage registry's collect-only pass never sees it.""" + test = type(self).test_every_test_carries_package_covers_and_source + request.applymarker(pytest.mark.covers("LOG-1", "LOG-2")) + assert result_properties(collected_item(request, test.__name__)) == ( + ("package", "root"), + ("covers", "LOG-1,LOG-2"), + ("source", f"tests/e2e/test_junit_properties.py:{test.__code__.co_firstlineno}"), + ) + + def test_attach_is_idempotent(self, request: pytest.FixtureRequest) -> None: + """Collection can run the hook more than once; a second pass must not + double the entries in the report.""" + item = collected_item(request, type(self).test_attach_is_idempotent.__name__) + attach_result_properties(item) + attach_result_properties(item) + assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] + + +class TestSuiteRoot: + def test_suite_root_names_this_file_s_real_home(self) -> None: + """SUITE_ROOT is hardcoded because the runner image has no repo to read it + from. Where there IS a checkout, prove the constant still points at us -- + otherwise a moved tests/e2e/ ships links that 404.""" + root = repo_root() + if root is None: + pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)") + assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve() + + +class TestDedupeCovers: + def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None: + assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C") diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index 9d918736262..bb33c90ddf3 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -29,6 +29,7 @@ export const E2E_PROXY_ADMIN_USER_ID = "e2e-proxy-admin"; export const E2E_PROXY_ADMIN_EMAIL = "admin@test.local"; export const E2E_INTERNAL_USER_ID = "e2e-internal-user"; export const E2E_INTERNAL_USER_EMAIL = "internal@test.local"; +export const E2E_TEAM_ADMIN_USER_ID = "e2e-team-admin"; // Key aliases for seeded test keys (match seed.sql) export const E2E_UPDATE_LIMITS_KEY_ALIAS = "e2eUpdateLimitsKey"; @@ -46,3 +47,5 @@ export const E2E_TEAM_ORG_ID = "e2e-team-org"; export const E2E_TEAM_ORG_ALIAS = "E2E Team In Org"; export const E2E_TEAM_NO_ADMIN_ID = "e2e-team-no-admin"; export const E2E_TEAM_NO_ADMIN_ALIAS = "E2E Team No Admin"; +export const E2E_TEAM_KEYGEN_ID = "e2e-team-keygen"; +export const E2E_TEAM_KEYGEN_ALIAS = "E2E Team Keygen"; diff --git a/tests/e2e/ui/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql index a1218633cdb..e77b4a16b3d 100644 --- a/tests/e2e/ui/fixtures/seed.sql +++ b/tests/e2e/ui/fixtures/seed.sql @@ -29,7 +29,7 @@ INSERT INTO "LiteLLM_UserTable" ("user_id", "user_email", "user_role", "teams", VALUES ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), + ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), @@ -63,6 +63,17 @@ INSERT INTO "LiteLLM_TeamTable" ( '[{"role":"user","user_id":"e2e-invitable-user"}]'::jsonb, '{}'::jsonb, '{"fake-openai-gpt-4"}', 0.0, '{}'::jsonb, '{}'::jsonb, false); +INSERT INTO "LiteLLM_TeamTable" ( + "team_id", "team_alias", "organization_id", "admins", "members", + "members_with_roles", "metadata", "models", "spend", "model_spend", "model_max_budget", "blocked", + "team_member_permissions" +) VALUES + ('e2e-team-keygen', 'E2E Team Keygen', NULL, + '{}', '{"e2e-internal-user"}', + '[{"role":"user","user_id":"e2e-internal-user"}]'::jsonb, + '{}'::jsonb, '{"fake-openai-gpt-4"}', 0.0, '{}'::jsonb, '{}'::jsonb, false, + '{"/key/generate"}'); + -- 6. Team Memberships (only user_id, team_id, spend — no created_at/updated_at) INSERT INTO "LiteLLM_TeamMembership" ("user_id", "team_id", "spend") VALUES @@ -72,6 +83,7 @@ VALUES ('e2e-removable-member', 'e2e-team-crud', 0.0), ('e2e-team-admin', 'e2e-team-delete', 0.0), ('e2e-internal-user', 'e2e-team-org', 0.0), + ('e2e-internal-user', 'e2e-team-keygen', 0.0), ('e2e-invitable-user', 'e2e-team-no-admin', 0.0); -- 7. Verification Tokens (API Keys) diff --git a/tests/e2e/ui/helpers/premium.ts b/tests/e2e/ui/helpers/premium.ts new file mode 100644 index 00000000000..28bc2e58bbc --- /dev/null +++ b/tests/e2e/ui/helpers/premium.ts @@ -0,0 +1,20 @@ +import * as fs from "fs"; +import { ADMIN_STORAGE_PATH } from "../constants"; + +/** + * Whether the proxy under test is licensed, read from the admin session JWT's `premium_user` + * claim. That is the same value the dashboard reads to enable premium-gated controls, so it + * describes the proxy Playwright is pointed at rather than the environment the runner happens + * to have, which are not the same machine when E2E_UI_BASE_URL points elsewhere. + */ +export function proxyIsPremium(): boolean { + const storage = JSON.parse(fs.readFileSync(ADMIN_STORAGE_PATH, "utf-8")) as { + cookies?: { name: string; value: string }[]; + }; + const token = storage.cookies?.find((cookie) => cookie.name === "token")?.value; + const payload = token?.split(".")[1]; + if (!payload) { + return false; + } + return JSON.parse(Buffer.from(payload, "base64url").toString("utf-8")).premium_user === true; +} diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index a2fc9463c94..25eb671fd0e 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -4,12 +4,16 @@ import { APIRequestContext, expect } from "@playwright/test"; export const CHAT_MODEL_A = "fake-openai-gpt-4"; export const CHAT_MODEL_B = "fake-anthropic-claude"; +/** The deployment each of those models routes to, as spend logs and usage breakdowns name it. */ +export const DEPLOYMENT_MODEL_A = "openai/fake-gpt-4"; +export const DEPLOYMENT_MODEL_B = "openai/fake-claude"; + /** The only completion text fixtures/mock_llm_server/server.py ever returns. */ export const MOCK_RESPONSE_TEXT = "This is a mock response."; export const masterKey = (): string => process.env.LITELLM_MASTER_KEY || "sk-1234"; -const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; +export const rootPath = (): string => process.env.SERVER_ROOT_PATH ?? ""; interface ChatOptions { model: string; @@ -84,15 +88,84 @@ export async function waitForSpendLog( throw new Error(`spend log for request ${requestId} never appeared (last /spend/logs status ${lastStatus})`); } +export async function waitForSpendLogByPrompt( + request: APIRequestContext, + prompt: string, + timeoutMs = 60_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let lastStatus = 0; + while (Date.now() < deadline) { + const res = await request.get(`${rootPath()}/spend/logs`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + lastStatus = res.status(); + if (res.ok()) { + const rows: { request_id?: string; messages?: unknown; proxy_server_request?: unknown }[] = await res.json(); + const row = (Array.isArray(rows) ? rows : []).find( + (candidate) => + JSON.stringify(candidate.messages ?? "").includes(prompt) || + JSON.stringify(candidate.proxy_server_request ?? "").includes(prompt), + ); + if (row?.request_id) { + return row.request_id; + } + } + await new Promise((r) => setTimeout(r, 2_000)); + } + throw new Error(`no spend log row carrying prompt ${prompt} appeared (last /spend/logs status ${lastStatus})`); +} + const isoDay = (d: Date): string => d.toISOString().slice(0, 10); +interface DailyActivityKey { + metrics?: { api_requests?: number }; +} + +interface DailyActivityPage { + results?: { breakdown?: { api_keys?: Record } }[]; + metadata?: { total_pages?: number }; +} + +const requestsOnPage = (body: DailyActivityPage, keyToken: string): number => + (body.results ?? []).reduce((sum, day) => sum + (day.breakdown?.api_keys?.[keyToken]?.metrics?.api_requests ?? 0), 0); + +/** + * The route paginates its per-key breakdown. Reading only the first page finds a key while the + * database is small and stops finding it once a run has generated more keys than one page holds, + * which reads as "the rollup is not running" when the rollup is fine. + */ +async function keyRequestsInDailyActivity( + request: APIRequestContext, + query: string, + keyToken: string, + page = 1, + seen = 0, +): Promise { + const res = await request.get(`${rootPath()}/user/daily/activity?${query}&page=${page}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + if (!res.ok()) { + return seen; + } + const body = (await res.json()) as DailyActivityPage; + const total = seen + requestsOnPage(body, keyToken); + return page >= (body.metadata?.total_pages ?? 1) + ? total + : keyRequestsInDailyActivity(request, query, keyToken, page + 1, total); +} + /** * The Usage page reads /user/daily/activity, a rollup written by a background job, and fetches it once * on mount. Navigating before the rollup lands leaves a stale render that never refreshes. + * + * The rollup lands request by request, so waiting only for the key to appear leaves a caller that + * sent several requests reading a partial count. Pass `minRequests` to wait for all of them. */ export async function waitForKeyInDailyActivity( request: APIRequestContext, keyToken: string, + minRequests = 1, timeoutMs = 120_000, ): Promise { const now = new Date(); @@ -101,25 +174,17 @@ export async function waitForKeyInDailyActivity( const query = `start_date=${isoDay(start)}&end_date=${isoDay(now)}`; const deadline = Date.now() + timeoutMs; - let lastStatus = 0; - while (Date.now() < deadline) { - const res = await request.get(`${rootPath()}/user/daily/activity?${query}`, { - headers: { Authorization: `Bearer ${masterKey()}` }, - }); - lastStatus = res.status(); - if (res.ok()) { - const body = await res.json(); - const seen = (body?.results ?? []).some( - (day: { breakdown?: { api_keys?: Record } }) => keyToken in (day.breakdown?.api_keys ?? {}), + for (;;) { + const seen = await keyRequestsInDailyActivity(request, query, keyToken); + if (seen >= minRequests) { + return; + } + if (Date.now() >= deadline) { + throw new Error( + `key ${keyToken} reached ${seen} of ${minRequests} requests in /user/daily/activity across every page; ` + + "the daily spend rollup may not be running", ); - if (seen) { - return; - } } await new Promise((r) => setTimeout(r, 3_000)); } - throw new Error( - `key ${keyToken} never appeared in /user/daily/activity (last status ${lastStatus}); ` + - "the daily spend rollup may not be running", - ); } diff --git a/tests/e2e/ui/tests/budgets/budgets.spec.ts b/tests/e2e/ui/tests/budgets/budgets.spec.ts new file mode 100644 index 00000000000..1ad1e488d25 --- /dev/null +++ b/tests/e2e/ui/tests/budgets/budgets.spec.ts @@ -0,0 +1,133 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { masterKey } from "../../helpers/traffic"; + +interface StoredBudget { + budget_id: string; + max_budget: number | null; + tpm_limit: number | null; + rpm_limit: number | null; + budget_duration: string | null; +} + +/** A different route from the one the table renders from, so a row that only lives in its cache fails here. */ +async function findBudget(page: PlaywrightPage, budgetId: string): Promise { + const res = await page.request.get("/budget/list", { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /budget/list (${res.status()})`).toBe(true); + return ((await res.json()) as StoredBudget[]).find((row) => row.budget_id === budgetId); +} + +async function createBudgetViaApi(page: PlaywrightPage, budget: Partial): Promise { + const res = await page.request.post("/budget/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: budget, + }); + expect(res.ok(), `POST /budget/new failed (${res.status()}): ${await res.text()}`).toBe(true); +} + +async function searchForBudget(page: PlaywrightPage, budgetId: string): Promise { + await page.getByPlaceholder("Search by budget ID").fill(budgetId); +} + +test.describe("Budgets", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Create a budget with rate limits and a spend cap", async ({ page }) => { + const budgetId = `e2e-budget-create-${Date.now()}`; + + await navigateToPage(page, Page.Budgets); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: "Create Budget" }).click(); + + const modal = page.getByRole("dialog", { name: "Create Budget" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + await modal.getByRole("textbox", { name: "Budget ID" }).fill(budgetId); + await modal.getByRole("spinbutton", { name: "Max Tokens per minute" }).fill("5000"); + await modal.getByRole("spinbutton", { name: "Max Requests per minute" }).fill("60"); + + await modal.getByRole("button", { name: "Optional Settings" }).click(); + await modal.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("25.5"); + await modal.getByRole("combobox", { name: "Reset Budget" }).click(); + await page.getByRole("option", { name: "weekly" }).click(); + + await modal.getByRole("button", { name: "Create Budget" }).click(); + await expect(modal).not.toBeVisible({ timeout: 10_000 }); + + await searchForBudget(page, budgetId); + const row = page.getByRole("row").filter({ hasText: budgetId }); + await expect(row).toBeVisible({ timeout: 10_000 }); + await expect(row).toContainText("$25.50"); + + const stored = await findBudget(page, budgetId); + expect(stored, `budget ${budgetId} readable from /budget/list`).toBeTruthy(); + expect(stored?.max_budget, "spend cap persisted").toBe(25.5); + expect(stored?.tpm_limit, "TPM limit persisted").toBe(5000); + expect(stored?.rpm_limit, "RPM limit persisted").toBe(60); + expect(stored?.budget_duration, "reset window persisted").toBe("7d"); + }); + + test("Raising a budget's spend cap leaves its rate limits alone", async ({ page }) => { + const budgetId = `e2e-budget-edit-${Date.now()}`; + await createBudgetViaApi(page, { budget_id: budgetId, max_budget: 10, tpm_limit: 1000, rpm_limit: 20 }); + + await navigateToPage(page, Page.Budgets); + await dismissFeedbackPopup(page); + + await searchForBudget(page, budgetId); + await expect(page.getByRole("row").filter({ hasText: budgetId })).toBeVisible({ timeout: 10_000 }); + + await page.getByTestId(`budget-actions-${budgetId}`).click(); + await page.getByTestId("budget-action-edit").click(); + + const modal = page.getByRole("dialog", { name: "Edit Budget" }); + await expect(modal).toBeVisible({ timeout: 10_000 }); + + await modal.getByRole("button", { name: "Optional Settings" }).click(); + await modal.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("99"); + await modal.getByRole("button", { name: "Save", exact: true }).click(); + await expect(modal).not.toBeVisible({ timeout: 10_000 }); + + await expect(page.getByRole("row").filter({ hasText: budgetId })).toContainText("$99.00", { timeout: 10_000 }); + + // Not hypothetical: the edit form posts the whole budget, so a field it fails to + // seed from the existing row goes to the server as null and silently clears. + const stored = await findBudget(page, budgetId); + expect(stored?.max_budget, "spend cap raised").toBe(99); + expect(stored?.tpm_limit, "TPM limit untouched by a spend-cap edit").toBe(1000); + expect(stored?.rpm_limit, "RPM limit untouched by a spend-cap edit").toBe(20); + }); + + test("Delete a budget", async ({ page }) => { + const budgetId = `e2e-budget-delete-${Date.now()}`; + await createBudgetViaApi(page, { budget_id: budgetId, max_budget: 5 }); + + await navigateToPage(page, Page.Budgets); + await dismissFeedbackPopup(page); + + await searchForBudget(page, budgetId); + await expect(page.getByRole("row").filter({ hasText: budgetId })).toBeVisible({ timeout: 10_000 }); + + await page.getByTestId(`budget-actions-${budgetId}`).click(); + await page.getByTestId("budget-action-delete").click(); + + const modal = page.getByRole("dialog", { name: "Delete Budget?" }); + await expect(modal).toBeVisible({ timeout: 5_000 }); + await modal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByRole("row").filter({ hasText: budgetId })).toHaveCount(0, { timeout: 10_000 }); + + // The row disappearing is a cache invalidation; the budget is gone when the route stops serving it. + await expect + .poll(async () => await findBudget(page, budgetId), { + message: `budget ${budgetId} still readable from /budget/list after delete`, + timeout: 15_000, + }) + .toBeUndefined(); + }); +}); diff --git a/tests/e2e/ui/tests/guardrails/guardrails.spec.ts b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts new file mode 100644 index 00000000000..1e43c7a2b22 --- /dev/null +++ b/tests/e2e/ui/tests/guardrails/guardrails.spec.ts @@ -0,0 +1,283 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_NO_ADMIN_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +interface StoredGuardrail { + guardrail_id: string; + guardrail_name: string | null; +} + +async function listGuardrails(page: PlaywrightPage): Promise { + const res = await page.request.get("/v2/guardrails/list", { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true); + return ((await res.json()) as { guardrails: StoredGuardrail[] }).guardrails; +} + +async function findGuardrail(page: PlaywrightPage, name: string): Promise { + return (await listGuardrails(page)).find((row) => row.guardrail_name === name); +} + +const createdGuardrails: string[] = []; + +async function createKeywordGuardrailViaApi(page: PlaywrightPage, name: string, keyword: string): Promise { + const res = await page.request.post("/guardrails", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + guardrail: { + guardrail_name: name, + litellm_params: { + guardrail: "litellm_content_filter", + mode: "pre_call", + default_on: false, + blocked_words: [{ keyword, action: "BLOCK" }], + }, + }, + }, + }); + expect(res.ok(), `POST /guardrails failed (${res.status()}): ${await res.text()}`).toBe(true); + createdGuardrails.push(name); + const guardrail = await findGuardrail(page, name); + expect(guardrail?.guardrail_id, `guardrail ${name} has an id`).toBeTruthy(); + return guardrail!.guardrail_id; +} + +async function openKeywordsStep(page: PlaywrightPage, name: string) { + await page.getByRole("button", { name: "Add New Guardrail" }).click(); + await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click(); + + const wizard = page.getByRole("dialog", { name: "Create guardrail" }); + await expect(wizard).toBeVisible({ timeout: 10_000 }); + + await wizard.getByRole("textbox", { name: "Guardrail Name" }).fill(name); + await wizard.getByRole("combobox", { name: "Guardrail Provider" }).click(); + // The content filter runs inside the proxy, so this is the one provider a test can + // configure end to end without standing up a third-party moderation service. + await page.getByRole("option", { name: /LiteLLM Content Filter/ }).click(); + + for (const step of ["Topics", "Patterns", "Keywords"]) { + await wizard.getByRole("button", { name: "Next" }).click(); + await expect(wizard).toContainText(step, { timeout: 10_000 }); + } + return wizard; +} + +test.describe("Guardrails", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test.afterEach(async ({ page }) => { + // Guardrails live in the database and show up in the table and the playground list, so a run + // that leaves them behind changes what the next run sees. + for (const name of createdGuardrails.splice(0)) { + const guardrail = await findGuardrail(page, name); + if (guardrail) { + const deleted = await page.request.delete(`/guardrails/${guardrail.guardrail_id}`, { + headers: { Authorization: `Bearer ${masterKey()}` }, + }); + expect(deleted.ok(), `DELETE /guardrails/${guardrail.guardrail_id} (${deleted.status()})`).toBe(true); + } + } + }); + + test("A guardrail created through the wizard blocks the keyword it was given", async ({ page }) => { + const stamp = Date.now(); + const guardrailName = `e2e-guardrail-create-${stamp}`; + // Unique per run so a concurrent test's prompt can never trip this guardrail, or vice versa. + const bannedKeyword = `e2ebanned${stamp}`; + + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + createdGuardrails.push(guardrailName); + const wizard = await openKeywordsStep(page, guardrailName); + + await wizard.getByRole("button", { name: "Add keyword" }).click(); + const keywordModal = page.getByRole("dialog", { name: "Add blocked keyword" }); + await expect(keywordModal).toBeVisible({ timeout: 10_000 }); + await keywordModal.getByPlaceholder("Enter sensitive keyword or phrase").fill(bannedKeyword); + await keywordModal.getByRole("button", { name: "Add", exact: true }).click(); + await expect(keywordModal).not.toBeVisible({ timeout: 10_000 }); + + await wizard.getByRole("button", { name: "Next" }).click(); + await wizard.getByRole("button", { name: "Create Guardrail" }).click(); + await expect(wizard).not.toBeVisible({ timeout: 15_000 }); + + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toBeVisible({ timeout: 15_000 }); + expect(await findGuardrail(page, guardrailName), "guardrail readable from /v2/guardrails/list").toBeTruthy(); + + // A row in the table only proves the record was written. The point of a guardrail is that it + // refuses traffic, so drive a request through it. + // + // Polled: a guardrail written through /guardrails reaches the request path on the proxy's + // periodic refresh, so the first call after creation can still be served unguarded. The + // assertion is unchanged, it just allows that refresh to land. + let blockedBody = ""; + await expect + .poll( + async () => { + const res = await page.request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }], + guardrails: [guardrailName], + }, + }); + blockedBody = await res.text(); + return res.status(); + }, + { message: "a prompt carrying the banned keyword is refused", timeout: 60_000 }, + ) + .toBe(400); + expect(blockedBody).toContain(bannedKeyword); + + const allowed = await page.request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: "hello there" }], + guardrails: [guardrailName], + }, + }); + expect(allowed.status(), "a clean prompt still gets through the same guardrail").toBe(200); + expect((await allowed.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + }); + + test("The Test Playground reports the verdict for the text it is given", async ({ page }) => { + const stamp = Date.now(); + const guardrailName = `e2e-guardrail-play-${stamp}`; + const bannedKeyword = `e2eplay${stamp}`; + await createKeywordGuardrailViaApi(page, guardrailName, bannedKeyword); + + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await page.getByRole("tab", { name: "Test Playground" }).click(); + // Every tab on this page stays mounted, so the other tabs' search boxes match too. + const playground = page.getByRole("tabpanel", { name: "Test Playground" }); + await playground.getByPlaceholder("Search guardrails...").fill(guardrailName); + await playground.getByText(guardrailName, { exact: true }).click(); + + const input = playground.getByPlaceholder("Enter text to test with guardrails..."); + await input.fill(`this sentence contains ${bannedKeyword}`); + await playground.getByRole("button", { name: /^Test 1 guardrail$/ }).click(); + + // The playground is where an admin checks a guardrail before rolling it out, so the + // verdict it prints has to be the one the gateway would give. + await expect(playground.getByText(`${guardrailName} - Error`)).toBeVisible({ timeout: 20_000 }); + await expect(playground.getByText(new RegExp(`Content blocked.*${bannedKeyword}`))).toBeVisible({ + timeout: 10_000, + }); + + await input.fill("this sentence is perfectly ordinary"); + await playground.getByRole("button", { name: /^Test 1 guardrail$/ }).click(); + + await expect(playground.getByText(`${guardrailName} - Error`)).toHaveCount(0, { timeout: 20_000 }); + await expect(playground.getByText("this sentence is perfectly ordinary").last()).toBeVisible({ timeout: 10_000 }); + }); + + test("Delete a guardrail", async ({ page }) => { + const stamp = Date.now(); + const guardrailName = `e2e-guardrail-delete-${stamp}`; + const guardrailId = await createKeywordGuardrailViaApi(page, guardrailName, `e2edelete${stamp}`); + + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toBeVisible({ timeout: 15_000 }); + + await page.getByTestId(`guardrail-actions-${guardrailId}`).click(); + await page.getByTestId("guardrail-action-delete").click(); + + const modal = page.getByRole("dialog"); + await expect(modal).toBeVisible({ timeout: 5_000 }); + await modal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(0, { timeout: 15_000 }); + + // The RC checklist deletes then reloads, because a row vanishing from the table has + // fooled us before; assert against the route the reload would read. + await expect + .poll(async () => await findGuardrail(page, guardrailName), { + message: `guardrail ${guardrailName} still listed after delete`, + timeout: 15_000, + }) + .toBeUndefined(); + }); + + test("Create a Presidio guardrail, see it in team settings, and delete it", async ({ page }) => { + const guardrailName = `e2e-presidio-${Date.now()}`; + + await navigateToPage(page, Page.Guardrails); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Add New Guardrail/i }).click(); + await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click(); + + const dialog = page.getByRole("dialog", { name: "Create guardrail" }); + await expect(dialog).toBeVisible({ timeout: 10_000 }); + + await dialog.getByLabel("Guardrail Name").fill(guardrailName); + + const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" }); + await providerSelect.click(); + await providerSelect.fill("Presidio"); + await page.getByRole("option", { name: "Presidio PII" }).click(); + + await dialog.getByLabel("Mode", { exact: true }).click(); + await page.keyboard.type("pre_call"); + await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 }); + await page.keyboard.press("Enter"); + await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 }); + await dialog.getByText("Create guardrail", { exact: true }).click(); + + await dialog.getByLabel("presidio_analyzer_api_base").fill("http://127.0.0.1:9999"); + await expect(dialog.getByLabel("presidio_analyzer_api_base")).toHaveValue("http://127.0.0.1:9999"); + await dialog.getByLabel("presidio_anonymizer_api_base").fill("http://127.0.0.1:9999"); + await expect(dialog.getByLabel("presidio_anonymizer_api_base")).toHaveValue("http://127.0.0.1:9999"); + + await dialog.getByRole("button", { name: "Next" }).click(); + await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 }); + await dialog.getByRole("button", { name: "Select All & Mask" }).click(); + + await dialog.getByRole("button", { name: "Create Guardrail" }).click(); + await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 }); + + const row = page.getByRole("row").filter({ hasText: guardrailName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + const guardrailsSelect = page.getByRole("combobox", { name: "Select guardrails" }); + await expect(guardrailsSelect).toBeVisible({ timeout: 10_000 }); + await guardrailsSelect.click(); + await guardrailsSelect.fill(guardrailName); + await expect(page.getByRole("option", { name: guardrailName })).toBeVisible({ timeout: 10_000 }); + await page.keyboard.press("Escape"); + + await navigateToPage(page, Page.Guardrails); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await row.getByRole("button", { name: "Open guardrail actions" }).click(); + await page.getByRole("menuitem", { name: "Delete" }).click(); + + const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" }); + await expect(deleteModal).toBeVisible({ timeout: 5_000 }); + await deleteModal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ + timeout: 10_000, + }); + await expect(row).toHaveCount(0, { timeout: 15_000 }); + + await page.reload(); + await expect(page.getByRole("button", { name: /Add New Guardrail/i })).toBeVisible({ timeout: 20_000 }); + await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(0); + }); +}); diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index 26e34dd2fe5..f392c5104da 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -3,10 +3,13 @@ import { E2E_INTERNAL_USER_KEY_ALIAS, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID, + E2E_TEAM_KEYGEN_ALIAS, INTERNAL_USER_STORAGE_PATH, } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage, clickTeamId } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; +import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; test.describe("Internal User", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -37,6 +40,55 @@ test.describe("Internal User", () => { await expect(page.getByRole("tab", { name: "Members" })).not.toBeVisible(); }); + test("Internal user creates a team key and uses it in the Playground", async ({ page, request }) => { + const suffix = Date.now(); + const auth = { Authorization: `Bearer ${masterKey()}` }; + + await navigateToPage(page, Page.ApiKeys); + + await page.getByRole("button", { name: /Create New Key/i }).click(); + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + await expect(page.getByRole("radio", { name: "You", exact: true })).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("radio", { name: "Another User" })).toHaveCount(0); + + const keyName = `e2e-internal-team-key-${suffix}`; + await page.getByLabel(/Key Name/).fill(keyName); + + const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); + await teamSelect.click(); + await page.keyboard.type(E2E_TEAM_KEYGEN_ALIAS); + await page.getByRole("option", { name: E2E_TEAM_KEYGEN_ALIAS }).first().click(); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Team Models", exact: true }).click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + await page.keyboard.press("Escape"); + + try { + await openPlayground(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + + const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(apiKey); + + await selectModel(page, CHAT_MODEL_A); + await sendMessage(page, `internal user team key ping ${keyName}`); + + await expect(page.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 60_000 }); + } finally { + await request.post("/key/delete", { headers: auth, data: { keys: [apiKey] } }); + } + }); + test("Virtual Keys page does not surface litellm-dashboard team keys", async ({ page }) => { await navigateToPage(page, Page.ApiKeys); diff --git a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts index 49e27a36673..62681e9ceb5 100644 --- a/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUserWithTeams.spec.ts @@ -1,14 +1,13 @@ import { test, expect } from "@playwright/test"; -import { INTERNAL_USER_STORAGE_PATH, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_ORG_ALIAS } from "../../constants"; +import { + INTERNAL_USER_STORAGE_PATH, + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_KEYGEN_ALIAS, + E2E_TEAM_ORG_ALIAS, +} from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage } from "../../helpers/navigation"; -/** - * Differential partner to internalUserNoTeam.spec.ts: the seeded - * e2e-internal-user belongs to exactly two teams, so the Create Key dropdown - * must list both. Without this, the no-team spec's "zero options" assertion - * would still pass against a bug that empties the dropdown for everyone. - */ test.describe("Internal User with team memberships", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -21,10 +20,9 @@ test.describe("Internal User with team memberships", () => { const teamSelect = page.getByTestId("team-dropdown").getByRole("combobox"); await teamSelect.click(); - // Both seeded memberships render, and nothing else does — proving the - // dropdown is scoped to the user's teams rather than empty or unfiltered. await expect(page.getByRole("option", { name: E2E_TEAM_CRUD_ALIAS })).toBeVisible({ timeout: 10_000 }); await expect(page.getByRole("option", { name: E2E_TEAM_ORG_ALIAS })).toBeVisible(); - await expect(page.getByRole("option")).toHaveCount(2); + await expect(page.getByRole("option", { name: E2E_TEAM_KEYGEN_ALIAS })).toBeVisible(); + await expect(page.getByRole("option")).toHaveCount(3); }); }); diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 610a88c6cd0..2748c91395f 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -2,7 +2,14 @@ import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwr import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; -import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; +import { + CHAT_MODEL_A, + MOCK_RESPONSE_TEXT, + sendChatCompletion, + waitForSpendLog, + waitForSpendLogByPrompt, +} from "../../helpers/traffic"; +import { openPlayground, selectModel, sendMessage } from "../../helpers/playground"; /** * Anchored to traffic this spec generates itself, with a unique prompt and end user per run, so it @@ -46,6 +53,23 @@ test.describe("Logs page", () => { permissions: ["clipboard-read", "clipboard-write"], }); + test("a chat sent from the Playground lands in Logs with its content", async ({ page, request }) => { + const prompt = `logs-playground-prompt-${uniqueSuffix()}`; + await openPlayground(page); + await selectModel(page, CHAT_MODEL_A); + await sendMessage(page, prompt); + await expect(page.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 60_000 }); + + const requestId = await waitForSpendLogByPrompt(request, prompt); + + const row = await openLogsForRequest(page, requestId); + await row.click(); + const drawer = page.getByRole("dialog").first(); + await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); + await expect(drawer.getByText(prompt, { exact: false }).first()).toBeVisible({ timeout: 20_000 }); + await expect(drawer.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 20_000 }); + }); + test("a served request expands to its request and response", async ({ page, request }) => { const prompt = `logs-detail-prompt-${uniqueSuffix()}`; const requestId = await sendChatCompletion(request, { diff --git a/tests/e2e/ui/tests/logs/logsFilters.spec.ts b/tests/e2e/ui/tests/logs/logsFilters.spec.ts new file mode 100644 index 00000000000..7174f339296 --- /dev/null +++ b/tests/e2e/ui/tests/logs/logsFilters.spec.ts @@ -0,0 +1,152 @@ +import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + createVirtualKey, + sendChatCompletion, + waitForSpendLog, +} from "../../helpers/traffic"; + +/** + * Every test mints its own key and asserts against request ids it generated, so a filter that + * quietly does nothing shows up as the other key's row still being on screen, and concurrent + * specs' traffic cannot decide the outcome. + */ + +const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */ +const requestLogsRows = (page: PlaywrightPage): Locator => + page.locator("table").filter({ visible: true }).first().locator("tbody tr"); + +const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true }); + +async function openLogs(page: PlaywrightPage): Promise { + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 }); +} + +async function openFilterDrawer(page: PlaywrightPage): Promise { + await visibleTestId(page, "datatable-filters-trigger").click(); + const drawer = page.getByRole("dialog", { name: "Filters" }); + await expect(drawer).toBeVisible({ timeout: 10_000 }); + return drawer; +} + +/** Picks a value in one of the drawer's searchable comboboxes and applies the filter. */ +async function applyComboboxFilter( + page: PlaywrightPage, + drawer: Locator, + comboboxLabel: string, + value: string, +): Promise { + await drawer.getByRole("combobox", { name: comboboxLabel }).click(); + await page.keyboard.type(value); + await page.getByRole("option", { name: value, exact: true }).first().click(); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); +} + +/** A request the key is not entitled to make, so the proxy refuses it and logs the refusal. */ +async function sendDeniedCompletion(request: APIRequestContext, apiKey: string): Promise { + const res = await request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + data: { model: CHAT_MODEL_B, messages: [{ role: "user", content: "denied" }] }, + }); + expect(res.status(), "a model outside the key's allow-list is refused").toBe(403); +} + +test.describe("Logs page filters", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("the Key Alias filter narrows the table to that key's requests", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const mine = await createVirtualKey(request, { key_alias: `e2e-logs-mine-${suffix}` }); + const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-theirs-${suffix}` }); + + const myRequestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-filter-mine-${suffix}`, + apiKey: mine.key, + }); + const theirRequestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-filter-theirs-${suffix}`, + apiKey: theirs.key, + }); + await waitForSpendLog(request, myRequestId); + await waitForSpendLog(request, theirRequestId); + + await openLogs(page); + const drawer = await openFilterDrawer(page); + await applyComboboxFilter(page, drawer, "Search a key alias", mine.alias!); + + await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 30_000 }); + // The filter is only doing its job if the other key's request is gone, not merely if ours is present. + await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(0, { timeout: 10_000 }); + }); + + test("the Status filter narrows the table to the refused request", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const alias = `e2e-logs-status-${suffix}`; + const scoped = await createVirtualKey(request, { key_alias: alias, models: [CHAT_MODEL_A] }); + + const servedRequestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-filter-served-${suffix}`, + apiKey: scoped.key, + }); + await sendDeniedCompletion(request, scoped.key); + await waitForSpendLog(request, servedRequestId); + + await openLogs(page); + const drawer = await openFilterDrawer(page); + await drawer.getByRole("combobox", { name: "Search a key alias" }).click(); + await page.keyboard.type(alias); + await page.getByRole("option", { name: alias, exact: true }).first().click(); + // The Status field labels its group, not the trigger, so it is addressed by the value it shows. + await drawer.getByRole("combobox").filter({ hasText: "All Statuses" }).click(); + await page.getByRole("option", { name: "Failure", exact: true }).click(); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); + + // Both requests were made by this key, so a Status filter that does nothing leaves the served one on screen. + await expect(requestLogsRows(page)).toHaveCount(1, { timeout: 30_000 }); + await expect(requestLogsRows(page)).toContainText("Failure"); + await expect(requestLogsRows(page).filter({ hasText: servedRequestId })).toHaveCount(0); + }); + + test("Reset Filters brings back the rows a filter hid", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const mine = await createVirtualKey(request, { key_alias: `e2e-logs-reset-mine-${suffix}` }); + const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-reset-theirs-${suffix}` }); + + const myRequestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-reset-mine-${suffix}`, + apiKey: mine.key, + }); + const theirRequestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-reset-theirs-${suffix}`, + apiKey: theirs.key, + }); + await waitForSpendLog(request, myRequestId); + await waitForSpendLog(request, theirRequestId); + + await openLogs(page); + const drawer = await openFilterDrawer(page); + await applyComboboxFilter(page, drawer, "Search a key alias", mine.alias!); + await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(0, { timeout: 30_000 }); + + // A filter you cannot clear is a page that looks empty forever, which is how it reads to a user. + await page.getByRole("button", { name: "Reset Filters" }).filter({ visible: true }).click(); + + await expect(requestLogsRows(page).filter({ hasText: theirRequestId })).toHaveCount(1, { timeout: 30_000 }); + await expect(requestLogsRows(page).filter({ hasText: myRequestId })).toHaveCount(1, { timeout: 10_000 }); + }); +}); diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 16ec94c1dc8..6877fc9c48d 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -1,7 +1,8 @@ -import { test, expect } from "@playwright/test"; +import { test, expect, type APIRequestContext } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; +import { masterKey } from "../../helpers/traffic"; test.describe("AI Hub (internal admin view)", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -77,4 +78,89 @@ test.describe("Public model hub (/ui/model_hub_table)", () => { // agents/MCP servers exist, so we don't assert on them in a fresh CI run. await expect(page.getByRole("tab", { name: "Model Hub" })).toBeVisible({ timeout: 10_000 }); }); + + test("Agent Hub and MCP Hub tabs render their public entries", async ({ page, request }) => { + const suffix = `${Date.now()}`; + const agentName = `e2e-public-agent-${suffix}`; + const mcpServerName = `e2e_public_mcp_${suffix}`; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const publicMcpServerIds = async (api: APIRequestContext): Promise => { + const res = await api.get("/public/mcp_hub"); + expect(res.ok(), `public mcp_hub read failed (${res.status()}): ${await res.text()}`).toBe(true); + const servers: { server_id: string }[] = await res.json(); + return servers.map((server) => server.server_id); + }; + + const seedPublicEntries = async ( + api: APIRequestContext, + priorMcpIds: string[], + ): Promise<{ agentId: string; serverId: string }> => { + const agentRes = await api.post("/v1/agents", { + headers: auth, + data: { + agent_name: agentName, + agent_card_params: { + name: agentName, + description: "E2E public agent", + version: "1.0.0", + url: "http://127.0.0.1:9999/", + capabilities: {}, + skills: [], + defaultInputModes: ["text"], + defaultOutputModes: ["text"], + }, + }, + }); + expect(agentRes.ok(), `agent create failed (${agentRes.status()}): ${await agentRes.text()}`).toBe(true); + const agentId = (await agentRes.json()).agent_id as string; + + const serverRes = await api.post("/v1/mcp/server", { + headers: auth, + data: { + server_name: mcpServerName, + url: "http://127.0.0.1:9999/mcp", + transport: "http", + description: "E2E public MCP server", + }, + }); + expect(serverRes.ok(), `mcp server create failed (${serverRes.status()}): ${await serverRes.text()}`).toBe(true); + const serverId = (await serverRes.json()).server_id as string; + + const agentPublicRes = await api.post(`/v1/agents/${agentId}/make_public`, { headers: auth }); + expect(agentPublicRes.ok(), `agent make_public failed: ${await agentPublicRes.text()}`).toBe(true); + const mcpPublicRes = await api.post("/v1/mcp/make_public", { + headers: auth, + data: { mcp_server_ids: [...priorMcpIds, serverId] }, + }); + expect(mcpPublicRes.ok(), `mcp make_public failed: ${await mcpPublicRes.text()}`).toBe(true); + + return { agentId, serverId }; + }; + + const priorMcpIds = await publicMcpServerIds(request); + const { agentId, serverId } = await seedPublicEntries(request, priorMcpIds); + try { + await page.goto(`/ui/model_hub_table?key=${masterKey()}`); + await dismissFeedbackPopup(page); + + const agentHubTab = page.getByRole("tab", { name: "Agent Hub" }); + await expect(agentHubTab).toBeVisible({ timeout: 15_000 }); + await agentHubTab.click(); + await expect(page.getByText("Available Agents")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("row").filter({ hasText: agentName })).toHaveCount(1, { timeout: 10_000 }); + await expect(page.getByText("E2E public agent").first()).toBeVisible(); + + const mcpHubTab = page.getByRole("tab", { name: "MCP Hub" }); + await expect(mcpHubTab).toBeVisible(); + await mcpHubTab.click(); + await expect(page.getByText("Available MCP Servers")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByRole("row").filter({ hasText: mcpServerName })).toHaveCount(1, { timeout: 10_000 }); + await expect(page.getByText("E2E public MCP server").first()).toBeVisible(); + } finally { + await request.post("/v1/mcp/make_public", { headers: auth, data: { mcp_server_ids: priorMcpIds } }); + await request.delete(`/v1/agents/${agentId}`, { headers: auth }); + await request.delete(`/v1/mcp/server/${serverId}`, { headers: auth }); + } + }); }); diff --git a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts index 84d1c01b452..de25ec1aac5 100644 --- a/tests/e2e/ui/tests/modelsPage/addModel.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/addModel.spec.ts @@ -5,6 +5,11 @@ import { navigateToPage } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; import { sendChatCompletion } from "../../helpers/traffic"; +import { proxyIsPremium } from "../../helpers/premium"; + +/** Four probes 13s apart span 39s, one PROXY_CONFIG_RELOAD_INTERVAL_SECONDS (30s) plus margin. */ +const CREDENTIAL_PROBE_SUCCESSES = 4; +const CREDENTIAL_PROBE_SPACING_MS = 13_000; /** The mock LLM as the proxy reaches it: same host locally, a sidecar in the deployed stack. */ const MOCK_LLM_BASE = `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`; @@ -35,7 +40,10 @@ async function selectProvider(page: PlaywrightPage, providerName: string) { const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); await providerDropdown.click(); await providerDropdown.fill(providerName); - await page.getByRole("option").filter({ hasText: exactly(providerName) }).click(); + await page + .getByRole("option") + .filter({ hasText: exactly(providerName) }) + .click(); await expect(providerDropdown).toHaveValue(providerName); } @@ -78,6 +86,9 @@ test.describe("Add Model", () => { }); test("Edit team model TPM and RPM limits", async ({ page }) => { + // /model/new refuses a team-scoped deployment on an unlicensed proxy, so there this fails in + // setup on a product gate rather than on a regression in the edit it covers. + test.skip(!proxyIsPremium(), "proxy under test is unlicensed — team-scoped models are premium"); const masterKey = users[Role.ProxyAdmin].password; const modelName = `e2e-team-model-${Date.now()}`; @@ -212,6 +223,120 @@ test.describe("Add Model", () => { .toBe(true); }); + test("Add a model with a stored credential, pass Test Connect, and serve traffic", async ({ page, request }) => { + const masterKey = users[Role.ProxyAdmin].password; + const auth = { Authorization: `Bearer ${masterKey}` }; + const credentialName = `e2e-cred-reuse-${Date.now()}`; + const createCred = await page.request.post("/credentials", { + headers: auth, + data: { + credential_name: credentialName, + credential_values: { api_key: "fake-key", api_base: MOCK_LLM_BASE }, + credential_info: { custom_llm_provider: "openai" }, + }, + }); + expect(createCred.ok(), `POST /credentials failed (${createCred.status()}): ${await createCred.text()}`).toBe(true); + + // The proxy's periodic credential refresh prunes its in-memory list against a database snapshot + // it took before this credential landed, so a credential that resolves right after POST + // /credentials can stop resolving until the refresh after that. Successes spanning a whole + // PROXY_CONFIG_RELOAD_INTERVAL_SECONDS prove it survived a refresh, after which it stays. + // Resolution fails open onto the ambient key, so losing it reads as a confusing upstream 404. + let consecutiveProbeSuccesses = 0; + await expect + .poll( + async () => { + const probe = await page.request.post("/health/test_connection", { + headers: auth, + data: { + litellm_params: { + model: "openai/fake-gpt-4", + custom_llm_provider: "openai", + litellm_credential_name: credentialName, + }, + model_info: {}, + mode: "chat", + }, + }); + const healthy = probe.ok() && (await probe.json()).status === "success"; + consecutiveProbeSuccesses = healthy ? consecutiveProbeSuccesses + 1 : 0; + return consecutiveProbeSuccesses; + }, + { + message: `stored credential ${credentialName} never stayed usable across a config reload`, + intervals: [0, CREDENTIAL_PROBE_SPACING_MS], + timeout: 110_000, + }, + ) + .toBeGreaterThanOrEqual(CREDENTIAL_PROBE_SUCCESSES); + + try { + await navigateToPage(page, Page.Models); + await page.getByRole("tab", { name: "Add Model" }).click(); + + await selectProvider(page, "OpenAI-Compatible Endpoints (Together AI, etc.)"); + + const publicName = `e2e-cred-model-${Date.now()}`; + uiAddedModelName = publicName; + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(publicName); + + const credentialSelect = page.getByRole("combobox", { name: "Existing Credentials" }); + await credentialSelect.click(); + await credentialSelect.fill(credentialName); + await page.getByRole("option", { name: credentialName, exact: true }).click(); + + await expect(page.locator("#api_key")).toHaveCount(0); + await expect(page.locator("#api_base")).toHaveCount(0); + + await page.getByRole("button", { name: "Test Connect" }).click(); + await expect(page.getByText("Connection Test Results")).toBeVisible({ timeout: 10_000 }); + await expect(page.getByTestId("connection-success-msg")).toBeVisible({ timeout: 30_000 }); + + const resultsModal = page.getByRole("dialog", { name: "Connection Test Results" }); + await resultsModal.locator('[data-slot="dialog-footer"]').getByRole("button", { name: "Close" }).click(); + await expect(resultsModal).toBeHidden({ timeout: 5_000 }); + + const created = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByRole("button", { name: "Add Model" }).last().click(); + }); + expect(created.litellm_params?.litellm_credential_name, "the picked credential goes on the wire").toBe( + credentialName, + ); + expect(created.litellm_params?.api_key, "no raw api key goes on the wire").toBeUndefined(); + + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + + await expect + .poll( + async () => { + try { + await sendChatCompletion(request, { model: publicName, prompt: `hello via ${credentialName}` }); + return true; + } catch { + return false; + } + }, + { + message: `model ${publicName} added with a stored credential never served a request`, + timeout: 30_000, + }, + ) + .toBe(true); + } finally { + const stored = uiAddedModelName ? await findDeploymentByName(page, uiAddedModelName) : undefined; + const id = stored?.model_info?.id; + if (id) { + await page.request.post("/model/delete", { headers: auth, data: { id } }); + uiAddedModelName = ""; + } + await page.request.delete(`/credentials/${credentialName}`, { headers: auth }); + } + }); + test("Test connection with bad credentials shows failure", async ({ page }) => { await navigateToPage(page, Page.Models); await page.getByRole("tab", { name: "Add Model" }).click(); @@ -348,7 +473,7 @@ test.describe("Add Model", () => { await page.waitForLoadState("networkidle"); await page.getByPlaceholder("Search model names").fill("cohere"); - + // Clearer failure than timing out on a row assertion when the table is empty. await expect(page.getByTestId("pagination-range")).toHaveText(/Showing \d+-\d+ of \d+/, { timeout: 15_000, @@ -356,10 +481,7 @@ test.describe("Add Model", () => { // Pin to one row carrying both the name and the team, so the sibling test's // team-less cohere row can't satisfy it. - const teamCohereRow = page - .getByRole("row") - .filter({ hasText: "cohere/" }) - .filter({ hasText: E2E_TEAM_CRUD_ID }); + const teamCohereRow = page.getByRole("row").filter({ hasText: "cohere/" }).filter({ hasText: E2E_TEAM_CRUD_ID }); await expect(teamCohereRow).toHaveCount(1, { timeout: 15_000 }); } finally { await deleteTeamScopedCohereModels(); diff --git a/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts new file mode 100644 index 00000000000..96abd9833c0 --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/deleteTeamModel.spec.ts @@ -0,0 +1,72 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; + +type DeploymentRow = { model_name?: string }; + +async function findDeploymentByName(page: PlaywrightPage, modelName: string): Promise { + const body = await readBack<{ data: DeploymentRow[] }>(page, "/v2/model/info"); + return body.data.find((row) => row.model_name === modelName); +} + +test.describe("Delete team model", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Delete a team-scoped model and verify it leaves the team's model list", async ({ page }) => { + const modelName = `e2e-team-model-delete-${Date.now()}`; + const createResponse = await page.request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + model_name: modelName, + litellm_params: { + model: "openai/fake-gpt-4", + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + }, + model_info: { team_id: E2E_TEAM_CRUD_ID }, + }, + }); + expect(createResponse.ok(), `/model/new failed: ${createResponse.status()} ${await createResponse.text()}`).toBe( + true, + ); + + await expect + .poll(async () => (await findDeploymentByName(page, modelName)) !== undefined, { + message: `deployment ${modelName} never appeared in /v2/model/info after create`, + timeout: 30_000, + }) + .toBe(true); + + await navigateToPage(page, Page.Models); + await page.getByPlaceholder("Search model names").fill(modelName); + + const row = page.getByRole("row").filter({ hasText: modelName }); + await expect(row).toHaveCount(1, { timeout: 15_000 }); + await expect(row.getByText(E2E_TEAM_CRUD_ID)).toBeVisible({ timeout: 10_000 }); + + await row.getByRole("button", { name: "Delete model" }).click(); + + const modal = page.getByRole("dialog", { name: "Delete Model" }); + await expect(modal).toBeVisible({ timeout: 5_000 }); + await expect(modal.getByText(modelName).first()).toBeVisible(); + await modal.getByRole("button", { name: "Delete", exact: true }).click(); + + await expect(page.getByText("Model deleted successfully").first()).toBeVisible({ timeout: 10_000 }); + await expect(row).toHaveCount(0, { timeout: 15_000 }); + + await expect + .poll(async () => await findDeploymentByName(page, modelName), { + message: `deployment ${modelName} still readable from /v2/model/info after delete`, + timeout: 15_000, + }) + .toBeUndefined(); + + await page.reload(); + await page.getByPlaceholder("Search model names").fill(modelName); + await expect(page.getByText("No models found").first()).toBeVisible({ timeout: 15_000 }); + await expect(page.getByRole("row").filter({ hasText: modelName })).toHaveCount(0); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts index f5bee68f245..deb7ae70d07 100644 --- a/tests/e2e/ui/tests/proxy-admin/keys.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keys.spec.ts @@ -1,15 +1,17 @@ import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH, - E2E_DELETE_KEY_ALIAS, E2E_REGENERATE_KEY_ALIAS, E2E_UPDATE_LIMITS_KEY_ALIAS, E2E_INTERNAL_USER_KEY_ALIAS, E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_CRUD_ID, } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; +import { proxyIsPremium } from "../../helpers/premium"; /** * Looks a key up by alias, undefined when none carries it. `return_full_object=true` is what makes @@ -23,6 +25,17 @@ async function findKeyByAlias(page: PlaywrightPage, alias: string): Promise row.key_alias === alias); } +/** A key this test owns, so deleting it costs the suite nothing on a retry or a second run. */ +async function createDeletableKey(page: PlaywrightPage): Promise { + const alias = `e2e-delete-key-${Date.now()}`; + const res = await page.request.post("/key/generate", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { key_alias: alias, team_id: E2E_TEAM_CRUD_ID }, + }); + expect(res.ok(), `POST /key/generate failed (${res.status()}): ${await res.text()}`).toBe(true); + return alias; +} + test.describe("Proxy Admin - Keys", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -67,6 +80,9 @@ test.describe("Proxy Admin - Keys", () => { }); test("Regenerate key", async ({ page }) => { + // The Regenerate Key button renders disabled when the proxy is unlicensed, so without one this + // fails on a product gate rather than on a regression. + test.skip(!proxyIsPremium(), "proxy under test is unlicensed — Regenerate Key is premium-gated"); await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); @@ -143,12 +159,16 @@ test.describe("Proxy Admin - Keys", () => { }); test("Delete key", async ({ page }) => { + // Deleting the seeded key leaves nothing for the next attempt, so the retries CI runs with are + // guaranteed to fail and the suite cannot run twice against one database. Bring our own. + const alias = await createDeletableKey(page); + await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - const keyRow = page.getByRole("row").filter({ hasText: E2E_DELETE_KEY_ALIAS }); + const keyRow = page.getByRole("row").filter({ hasText: alias }); await expect(keyRow).toBeVisible({ timeout: 10_000 }); - await keyRow.getByRole("button", { name: E2E_DELETE_KEY_ALIAS }).click(); + await keyRow.getByRole("button", { name: alias }).click(); await expect(page.getByText("Back to Keys")).toBeVisible({ timeout: 10_000 }); @@ -157,7 +177,7 @@ test.describe("Proxy Admin - Keys", () => { const modal = page.getByRole("dialog", { name: "Delete Key" }); await expect(modal).toBeVisible({ timeout: 5_000 }); - await modal.locator("input").fill(E2E_DELETE_KEY_ALIAS); + await modal.locator("input").fill(alias); const deleteButton = modal.getByRole("button", { name: "Delete", exact: true }); await expect(deleteButton).toBeEnabled(); @@ -167,8 +187,8 @@ test.describe("Proxy Admin - Keys", () => { // The key is gone when the management API stops returning it, not when the toast says so. await expect - .poll(async () => await findKeyByAlias(page, E2E_DELETE_KEY_ALIAS), { - message: `key ${E2E_DELETE_KEY_ALIAS} still readable from /key/list after delete`, + .poll(async () => await findKeyByAlias(page, alias), { + message: `key ${alias} still readable from /key/list after delete`, timeout: 15_000, }) .toBeUndefined(); diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts new file mode 100644 index 00000000000..5a8bc84cc13 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -0,0 +1,97 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +test.describe("Second proxy admin", () => { + test.use({ storageState: { cookies: [], origins: [] } }); + + test("an invited admin can log in, mint a key, and call a model with it", async ({ page, browser, request }) => { + const suffix = Date.now(); + const email = `second-admin-${suffix}@test.local`; + const password = "e2e-second-admin-password"; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const inviteAdminUser = async (): Promise => { + const adminContext = await browser.newContext({ storageState: ADMIN_STORAGE_PATH }); + try { + const adminPage = await adminContext.newPage(); + await navigateToPage(adminPage, Page.Users); + await dismissFeedbackPopup(adminPage); + + await adminPage.getByRole("button", { name: "+ Invite User", exact: true }).click(); + const dialog = adminPage.getByRole("dialog", { name: "Invite User" }); + await expect(dialog).toBeVisible({ timeout: 5_000 }); + + await dialog.getByLabel("User Email").fill(email); + + await dialog.getByLabel(/Global Proxy Role/).click(); + await adminPage.getByRole("option", { name: /Admin \(All Permissions\)/ }).click(); + + const createdResponse = adminPage.waitForResponse( + (res) => res.url().includes("/user/new") && res.request().method() === "POST", + ); + await dialog.getByRole("button", { name: "Invite User" }).click(); + const createdBody = await (await createdResponse).json(); + const createdUserId = (createdBody.data?.user_id ?? createdBody.user_id) as string; + expect(createdUserId, "created user id from /user/new").toBeTruthy(); + + await expect(adminPage.getByText("API user Created").first()).toBeVisible({ timeout: 10_000 }); + return createdUserId; + } finally { + await adminContext.close(); + } + }; + + const userId = await inviteAdminUser(); + try { + const passwordRes = await request.post("/user/update", { + headers: auth, + data: { user_email: email, password }, + }); + expect(passwordRes.ok(), `setting password failed (${passwordRes.status()}): ${await passwordRes.text()}`).toBe( + true, + ); + + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(password); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 }); + await dismissFeedbackPopup(page); + + await navigateToPage(page, Page.ApiKeys); + await page.getByRole("button", { name: /Create New Key/i }).click(); + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + await page.getByLabel(/Key Name/).fill(`e2e-second-admin-key-${suffix}`); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "All Proxy Models", exact: true }).click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + const apiKey = (await page.getByRole("dialog", { name: "Save your Key" }).locator("pre").innerText()).trim(); + expect(apiKey).toMatch(/^sk-/); + await page.keyboard.press("Escape"); + + const response = await page.request.post("/chat/completions", { + headers: { Authorization: `Bearer ${apiKey}` }, + data: { + model: CHAT_MODEL_A, + messages: [{ role: "user", content: `second admin ping ${suffix}` }], + }, + }); + expect(response.status()).toBe(200); + const body = await response.json(); + expect(body.choices?.[0]?.message?.content).toBe(MOCK_RESPONSE_TEXT); + } finally { + if (userId) { + await request.post("/user/delete", { headers: auth, data: { user_ids: [userId] } }); + } + } + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teamSettings.spec.ts b/tests/e2e/ui/tests/proxy-admin/teamSettings.spec.ts new file mode 100644 index 00000000000..e71945a4ccd --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/teamSettings.spec.ts @@ -0,0 +1,162 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; + +interface TeamInfo { + team_id: string; + team_alias: string; + models: string[]; + max_budget: number | null; + tpm_limit: number | null; + rpm_limit: number | null; + metadata: Record | null; + members_with_roles: { user_id?: string; role?: string }[]; +} + +/** + * Each test owns a team it created, rather than editing a seeded one, so a save that clobbers a + * field cannot take another spec's fixture down with it. + */ +async function createTeam(page: PlaywrightPage, alias: string, members: string[] = []): Promise { + const res = await page.request.post("/team/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { + team_alias: alias, + models: [CHAT_MODEL_A], + members_with_roles: members.map((user_id) => ({ user_id, role: "user" })), + }, + }); + expect(res.ok(), `POST /team/new failed (${res.status()}): ${await res.text()}`).toBe(true); + return (await res.json()).team_id as string; +} + +/** + * A member of this test's own, not one of the seeded users. Putting a seeded user on an extra team + * changes what every spec that asserts on their memberships sees. + */ +async function createMember(page: PlaywrightPage, userId: string): Promise { + const res = await page.request.post("/user/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { user_id: userId, user_role: "internal_user", auto_create_key: false }, + }); + expect(res.ok(), `POST /user/new failed (${res.status()}): ${await res.text()}`).toBe(true); + return userId; +} + +async function teamInfo(page: PlaywrightPage, teamId: string): Promise { + const body = await readBack<{ team_info: TeamInfo }>(page, `/team/info?team_id=${encodeURIComponent(teamId)}`); + return body.team_info; +} + +async function openTeamSettings(page: PlaywrightPage, teamId: string): Promise { + await navigateToPage(page, Page.Teams); + await dismissFeedbackPopup(page); + await clickTeamId(page, teamId); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByRole("button", { name: "Save Changes" })).toBeVisible({ timeout: 10_000 }); +} + +test.describe("Proxy Admin - Team settings", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Setting a team's spend cap and rate limits leaves its models and members alone", async ({ page }) => { + const stamp = Date.now(); + const alias = `e2e-team-limits-${stamp}`; + const member = await createMember(page, `e2e-team-limits-member-${stamp}`); + const teamId = await createTeam(page, alias, [member]); + const before = await teamInfo(page, teamId); + + await openTeamSettings(page, teamId); + + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("42.5"); + await page.getByRole("spinbutton", { name: "Tokens per minute Limit (TPM)" }).fill("7000"); + await page.getByRole("spinbutton", { name: "Requests per minute Limit (RPM)" }).fill("70"); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll( + async () => { + const team = await teamInfo(page, teamId); + return [team.max_budget, team.tpm_limit, team.rpm_limit]; + }, + { message: "team limits did not persist", timeout: 20_000 }, + ) + .toEqual([42.5, 7000, 70]); + + // The Settings form posts the whole team. A field it fails to seed goes back as null, and + // the toast still says success, so pin the fields this edit had no business touching. + const after = await teamInfo(page, teamId); + expect(after.models, "model access untouched by a limits edit").toEqual(before.models); + expect( + after.members_with_roles.map((member) => member.user_id).sort(), + "membership untouched by a limits edit", + ).toEqual(before.members_with_roles.map((member) => member.user_id).sort()); + }); + + test("A model alias added on the Settings tab serves traffic under the alias name", async ({ page }) => { + const stamp = Date.now(); + const alias = `e2e-team-alias-${stamp}`; + const modelAlias = `e2e-alias-${stamp}`; + const teamId = await createTeam(page, alias); + + await openTeamSettings(page, teamId); + + await page.getByRole("textbox", { name: "Alias Name" }).fill(modelAlias); + await page.getByRole("combobox", { name: "Select target model" }).click(); + await page.getByRole("option", { name: CHAT_MODEL_A, exact: true }).first().click(); + await page.getByRole("button", { name: "Add Alias" }).click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await teamInfo(page, teamId)).models, { message: "team lost its models", timeout: 20_000 }) + .toEqual([CHAT_MODEL_A]); + + const keyRes = await page.request.post("/key/generate", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { team_id: teamId, key_alias: `e2e-alias-key-${stamp}` }, + }); + expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()})`).toBe(true); + const teamKey = (await keyRes.json()).key as string; + + // An alias the team can see but cannot call is the actual complaint; the readback alone + // would pass for an alias the router never resolves. + const served = await page.request.post("/v1/chat/completions", { + headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" }, + data: { model: modelAlias, messages: [{ role: "user", content: "ping" }] }, + }); + expect(served.status(), `a team key calling ${modelAlias} is served`).toBe(200); + expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + }); + + test("Team metadata added as key-value pairs survives a reload", async ({ page }) => { + const stamp = Date.now(); + const alias = `e2e-team-metadata-${stamp}`; + const metadataValue = `cost-center-${stamp}`; + const teamId = await createTeam(page, alias); + + await openTeamSettings(page, teamId); + + await page.getByRole("button", { name: "Add Key-Value Pair" }).click(); + await page.getByPlaceholder("Key", { exact: true }).last().fill("owner"); + await page.getByPlaceholder("Value", { exact: true }).last().fill(metadataValue); + await page.getByRole("button", { name: "Save Changes" }).click(); + + await expect + .poll(async () => (await teamInfo(page, teamId)).metadata?.owner, { + message: "team metadata did not persist", + timeout: 20_000, + }) + .toBe(metadataValue); + + // Reopening the form is the step that catches metadata the page writes but cannot read back. + await page.reload(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByPlaceholder("Key", { exact: true })).toHaveValue("owner", { timeout: 15_000 }); + await expect(page.getByPlaceholder("Value", { exact: true })).toHaveValue(metadataValue); + }); +}); diff --git a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts index 7383b452162..303e4488e09 100644 --- a/tests/e2e/ui/tests/proxy-admin/teams.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/teams.spec.ts @@ -1,14 +1,9 @@ import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; -import { - ADMIN_STORAGE_PATH, - E2E_TEAM_CRUD_ID, - E2E_TEAM_DELETE_ALIAS, - E2E_TEAM_NO_ADMIN_ID, - E2E_TEAM_ORG_ID, -} from "../../constants"; +import { ADMIN_STORAGE_PATH, E2E_TEAM_CRUD_ID, E2E_TEAM_NO_ADMIN_ID, E2E_TEAM_ORG_ID } from "../../constants"; import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; import { readBack } from "../../helpers/roundTrip"; +import { masterKey } from "../../helpers/traffic"; /** GET /team/list returns a bare array of teams, each carrying team_alias/team_id. */ async function findTeamByAlias(page: PlaywrightPage, alias: string): Promise | undefined> { @@ -25,6 +20,17 @@ async function teamMemberEmails(page: PlaywrightPage, teamId: string): Promise member.user_email ?? "").filter(Boolean); } +/** A team this test owns, so deleting it costs the suite nothing on a retry or a second run. */ +async function createDeletableTeam(page: PlaywrightPage): Promise { + const alias = `e2e-delete-team-${Date.now()}`; + const res = await page.request.post("/team/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { team_alias: alias, models: ["fake-openai-gpt-4"] }, + }); + expect(res.ok(), `POST /team/new failed (${res.status()}): ${await res.text()}`).toBe(true); + return alias; +} + test.describe("Proxy Admin - Teams", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); @@ -121,10 +127,14 @@ test.describe("Proxy Admin - Teams", () => { }); test("Delete a team", async ({ page }) => { + // Deleting the seeded team leaves nothing for the next attempt, so the retries CI runs with are + // guaranteed to fail and the suite cannot run twice against one database. Bring our own. + const alias = await createDeletableTeam(page); + await navigateToPage(page, Page.Teams); await dismissFeedbackPopup(page); - const teamRow = page.locator("tr", { hasText: E2E_TEAM_DELETE_ALIAS }).first(); + const teamRow = page.locator("tr", { hasText: alias }).first(); await expect(teamRow).toBeVisible({ timeout: 10_000 }); // Actions live in a kebab menu: open it, then click "Delete team". await teamRow.locator('[data-testid^="team-actions-"]').click(); @@ -132,15 +142,15 @@ test.describe("Proxy Admin - Teams", () => { const modal = page.getByRole("dialog", { name: "Delete Team?" }); await expect(modal).toBeVisible({ timeout: 5_000 }); - await modal.locator("input").fill(E2E_TEAM_DELETE_ALIAS); + await modal.locator("input").fill(alias); await modal.getByRole("button", { name: /Force Delete|Delete/i }).click(); await expect(teamRow).not.toBeVisible({ timeout: 10_000 }); // A row vanishing is local state, which happens whether or not the delete landed. await expect - .poll(async () => await findTeamByAlias(page, E2E_TEAM_DELETE_ALIAS), { - message: `team ${E2E_TEAM_DELETE_ALIAS} still readable from /team/list after delete`, + .poll(async () => await findTeamByAlias(page, alias), { + message: `team ${alias} still readable from /team/list after delete`, timeout: 15_000, }) .toBeUndefined(); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index 5116cb5df19..f3c031f0172 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -1,6 +1,7 @@ import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; import { E2E_INTERNAL_USER_KEY_ALIAS, + E2E_TEAM_ADMIN_USER_ID, E2E_TEAM_CRUD_ALIAS, E2E_TEAM_CRUD_ID, TEAM_ADMIN_STORAGE_PATH, @@ -8,6 +9,8 @@ import { import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; +import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; +import { keySourceSelect, modelSelect, onlyVisible, openPlayground } from "../../helpers/playground"; /** * Every identifier a roster is addressable by. Which of user_id / user_email is populated depends on @@ -32,7 +35,44 @@ async function findKeyByAlias(page: PlaywrightPage, alias: string): Promise row.key_alias === alias); } +/** A member this test adds itself, so removing it costs the suite nothing on a retry or a re-run. */ +async function addRemovableMember(page: PlaywrightPage, registerForCleanup: string[]): Promise { + const userId = `e2e-removable-${Date.now()}`; + // Claimed before the call: /user/new can persist the user and still answer non-2xx, and the id is + // ours either way, so registering it up front is what no failure path can skip. + registerForCleanup.push(userId); + const created = await page.request.post("/user/new", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { user_id: userId, user_role: "internal_user", auto_create_key: false }, + }); + expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true); + + const added = await page.request.post("/team/member_add", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { team_id: E2E_TEAM_CRUD_ID, member: { user_id: userId, role: "user" } }, + }); + expect(added.ok(), `POST /team/member_add failed (${added.status()}): ${await added.text()}`).toBe(true); + return userId; +} + test.describe("Team Admin", () => { + const createdMembers: string[] = []; + + test.afterEach(async ({ page }) => { + // Runs on the failure path too, which a call at the end of the test body would not. Ids are + // claimed before the user is created, so the delete is attempted unconditionally and only its + // own 404 counts as never persisted; any other answer is a cleanup failure worth reporting + // rather than a reason to leave the user behind. + for (const userId of createdMembers.splice(0)) { + const deleted = await page.request.post("/user/delete", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { user_ids: [userId] }, + }); + const settled = deleted.ok() || deleted.status() === 404; + expect(settled, `POST /user/delete for ${userId} (${deleted.status()}): ${await deleted.text()}`).toBe(true); + } + }); + test.use({ storageState: TEAM_ADMIN_STORAGE_PATH }); test("Team admin can see all team keys including internal user keys", async ({ page }) => { @@ -92,6 +132,10 @@ test.describe("Team Admin", () => { }); test("Team admin can remove a member from their team", async ({ page }) => { + // Removing the seeded member leaves nothing for the next attempt, so the retries CI runs with + // are guaranteed to fail and the suite cannot run twice against one database. Bring our own. + const memberId = await addRemovableMember(page, createdMembers); + await navigateToPage(page, Page.Teams); await dismissFeedbackPopup(page); @@ -99,9 +143,9 @@ test.describe("Team Admin", () => { await page.getByRole("tab", { name: "Members" }).click(); - // Seeded members appear in the roster by user_id (members_with_roles has no - // email), so match the row on the user_id rather than the email. - const row = page.locator("tr", { hasText: "e2e-removable-member" }).first(); + // Members appear in the roster by user_id (members_with_roles has no email), so match + // the row on the user_id rather than the email. + const row = page.locator("tr", { hasText: memberId }).first(); await expect(row).toBeVisible({ timeout: 10_000 }); await row.getByTestId("delete-member").click(); @@ -114,7 +158,7 @@ test.describe("Team Admin", () => { // Removing the wrong member is exactly what a success toast hides, so pin both halves. expect(remove.team_id, "delete targets the team being viewed").toBe(E2E_TEAM_CRUD_ID); expect([remove.user_id, remove.user_email], "delete identifies the member whose row was clicked").toContain( - "e2e-removable-member", + memberId, ); await expect(page.getByText("Team member removed successfully").first()).toBeVisible({ timeout: 10_000 }); @@ -125,7 +169,92 @@ test.describe("Team Admin", () => { message: "removed member is still on the team", timeout: 15_000, }) - .not.toContain("e2e-removable-member"); + .not.toContain(memberId); + }); + + test("Team admin sees all team models in the Playground model dropdown", async ({ page, request }) => { + const suffix = Date.now(); + const teamModelName = `e2e-team-dropdown-model-${suffix}`; + const auth = { Authorization: `Bearer ${masterKey()}` }; + + const teamRes = await request.post("/team/new", { + headers: auth, + data: { + team_alias: `e2e-playground-team-${suffix}`, + models: [CHAT_MODEL_A], + members_with_roles: [{ role: "admin", user_id: E2E_TEAM_ADMIN_USER_ID }], + }, + }); + expect(teamRes.ok(), `team create failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true); + const teamId = (await teamRes.json()).team_id as string; + + try { + const modelRes = await request.post("/model/new", { + headers: auth, + data: { + model_name: teamModelName, + litellm_params: { + model: "openai/fake-gpt-4", + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + }, + model_info: { team_id: teamId }, + }, + }); + expect(modelRes.ok(), `model create failed (${modelRes.status()}): ${await modelRes.text()}`).toBe(true); + const modelId = (await modelRes.json()).model_info?.id as string; + + try { + const keyRes = await request.post("/key/generate", { headers: auth, data: { team_id: teamId } }); + expect(keyRes.ok(), `key generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true); + const teamKey = (await keyRes.json()).key as string; + + try { + await expect + .poll( + async () => { + const res = await request.get("/model_group/info", { + headers: { Authorization: `Bearer ${teamKey}` }, + }); + if (!res.ok()) return false; + const body: { data?: { model_group?: string }[] } = await res.json(); + return (body.data ?? []).some((group) => group.model_group === teamModelName); + }, + { + message: `model group ${teamModelName} never became visible to the team key`, + timeout: 30_000, + }, + ) + .toBe(true); + + await openPlayground(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + + const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(teamKey); + + const select = modelSelect(page); + await select.click(); + await select.fill(teamModelName); + await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ + timeout: 15_000, + }); + + await select.fill(CHAT_MODEL_A); + await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ + timeout: 15_000, + }); + } finally { + await request.post("/key/delete", { headers: auth, data: { keys: [teamKey] } }); + } + } finally { + await request.post("/model/delete", { headers: auth, data: { id: modelId } }); + } + } finally { + await request.post("/team/delete", { headers: auth, data: { team_ids: [teamId] } }); + } }); test("Team admin can create a team key with All Team Models", async ({ page }) => { diff --git a/tests/e2e/ui/tests/usage/usageActivityTabs.spec.ts b/tests/e2e/ui/tests/usage/usageActivityTabs.spec.ts new file mode 100644 index 00000000000..2ee5ae3e392 --- /dev/null +++ b/tests/e2e/ui/tests/usage/usageActivityTabs.spec.ts @@ -0,0 +1,135 @@ +import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { + CHAT_MODEL_A, + CHAT_MODEL_B, + DEPLOYMENT_MODEL_A, + DEPLOYMENT_MODEL_B, + createVirtualKey, + masterKey, + rootPath, + sendChatCompletion, + waitForKeyInDailyActivity, + waitForSpendLog, +} from "../../helpers/traffic"; + +/** + * Covers the per-entity breakdowns on /ui/usage. The page-level totals move with every other spec's + * traffic, so each assertion is scoped to a key this test minted and to the requests it sent. + */ + +/** Each breakdown renders one expandable card per entity, named " $x.xx N requests". */ +const entityCard = (page: PlaywrightPage, tab: string, name: string): Locator => + page.getByRole("tabpanel", { name: tab }).getByRole("button", { name: new RegExp(`^${name}\\s`) }); + +async function openUsageTab(page: PlaywrightPage, tab: string): Promise { + await navigateToPage(page, Page.NewUsage); + await dismissFeedbackPopup(page); + await page.getByRole("tab", { name: tab }).click(); + const panel = page.getByRole("tabpanel", { name: tab }); + await expect(panel).toBeVisible({ timeout: 30_000 }); + return panel; +} + +/** Sends `count` completions on one model and waits for each to reach the spend log. */ +async function sendTraffic( + request: Parameters[0], + apiKey: string, + model: string, + count: number, + label: string, +): Promise { + for (let i = 0; i < count; i++) { + const requestId = await sendChatCompletion(request, { model, prompt: `${label} ${i}`, apiKey }); + await waitForSpendLog(request, requestId); + } +} + +test.describe("Usage page activity tabs", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Key Activity breaks a key's traffic down by model", async ({ page, request }) => { + const alias = `e2e-usage-keyact-${Date.now()}`; + const { key, token } = await createVirtualKey(request, { key_alias: alias }); + + // An uneven split, so a breakdown that lumps everything into one row or attributes to the + // wrong model cannot land on these numbers by accident. + await sendTraffic(request, key, CHAT_MODEL_A, 2, alias); + await sendTraffic(request, key, CHAT_MODEL_B, 1, alias); + await waitForKeyInDailyActivity(request, token, 3); + + await openUsageTab(page, "Key Activity"); + + const card = entityCard(page, "Key Activity", alias); + await expect(card, `${alias} missing from Key Activity`).toBeVisible({ timeout: 30_000 }); + await expect(card).toContainText("3 requests"); + + // Every key gets a card, and the page opens the first one. Scope to this key's own section, + // which the collapsible renders as the trigger's next sibling. + await card.click(); + const details = card.locator("xpath=following-sibling::*[1]"); + const successfulFor = (model: string) => + details.getByRole("row").filter({ hasText: model }).getByRole("cell").nth(2); // Model | Spend | Successful | Failed | Tokens + + await expect(successfulFor(DEPLOYMENT_MODEL_A)).toHaveText("2", { timeout: 20_000 }); + await expect(successfulFor(DEPLOYMENT_MODEL_B)).toHaveText("1"); + }); + + test("Model Activity can name its models by deployment instead of by public name", async ({ page, request }) => { + const alias = `e2e-usage-modelact-${Date.now()}`; + const { key, token } = await createVirtualKey(request, { key_alias: alias }); + await sendTraffic(request, key, CHAT_MODEL_A, 1, alias); + await waitForKeyInDailyActivity(request, token); + + const panel = await openUsageTab(page, "Model Activity"); + + await expect(entityCard(page, "Model Activity", CHAT_MODEL_A), `${CHAT_MODEL_A} missing`).toBeVisible({ + timeout: 30_000, + }); + // Nothing is published under the deployment's name, so its absence here is what makes the + // toggle below a real change of key rather than a relabelled button. + await expect(entityCard(page, "Model Activity", DEPLOYMENT_MODEL_A)).toHaveCount(0); + + // Admins reconcile provider bills against the deployment, not the name their users call. + await panel.getByRole("button", { name: "Litellm Model Name" }).click(); + await expect(entityCard(page, "Model Activity", DEPLOYMENT_MODEL_A)).toBeVisible({ timeout: 20_000 }); + }); + + test("Filter by user narrows Key Activity to that user's keys", async ({ page, request }) => { + const stamp = Date.now(); + const email = `e2e-usage-owner-${stamp}@test.local`; + const ownedAlias = `e2e-usage-owned-${stamp}`; + const otherAlias = `e2e-usage-other-${stamp}`; + + const userRes = await request.post(`${rootPath()}/user/new`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { user_email: email, user_role: "internal_user", auto_create_key: false }, + }); + expect(userRes.ok(), `POST /user/new failed (${userRes.status()})`).toBe(true); + const userId = (await userRes.json()).user_id as string; + + const owned = await createVirtualKey(request, { key_alias: ownedAlias, user_id: userId }); + const other = await createVirtualKey(request, { key_alias: otherAlias }); + await sendTraffic(request, owned.key, CHAT_MODEL_A, 1, ownedAlias); + await sendTraffic(request, other.key, CHAT_MODEL_A, 1, otherAlias); + await waitForKeyInDailyActivity(request, owned.token); + await waitForKeyInDailyActivity(request, other.token); + + await openUsageTab(page, "Key Activity"); + await expect(entityCard(page, "Key Activity", otherAlias)).toBeVisible({ timeout: 30_000 }); + + await page.getByRole("combobox", { name: "Search users by email" }).click(); + await page.keyboard.type(email); + await page + .getByRole("option", { name: new RegExp(email) }) + .first() + .click(); + + // The filter earns its place only by dropping the other key; the owned key showing up + // proves nothing on a page that already listed every key. + await expect(entityCard(page, "Key Activity", otherAlias)).toHaveCount(0, { timeout: 30_000 }); + await expect(entityCard(page, "Key Activity", ownedAlias)).toBeVisible({ timeout: 20_000 }); + }); +}); diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index f61ab018e1b..3d057cfa2c9 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -1,10 +1,11 @@ -import { test, expect, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; import { - CHAT_MODEL_A, createVirtualKey, + masterKey, + rootPath, sendChatCompletion, waitForKeyInDailyActivity, waitForSpendLog, @@ -27,9 +28,105 @@ async function openUsage(page: PlaywrightPage): Promise { return card; } +/** The upstream fixtures/config.yml points its models at, so the mock server answers this too. */ +const MOCK_DEPLOYMENT = "openai/fake-gpt-4"; + +/** A deployment whose traffic costs real money, so the key that used it outranks the $0 crowd. */ +async function createPricedDeployment( + request: APIRequestContext, + label: string, + registerForCleanup: string[], +): Promise<{ modelName: string }> { + const modelName = `e2e-usage-priced-${label}`; + // Claimed before the call: /model/new can persist the deployment and still answer non-2xx, so a + // name recorded up front is the only registration no response shape can skip. + registerForCleanup.push(modelName); + const res = await request.post("/model/new", { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { + model_name: modelName, + litellm_params: { + model: MOCK_DEPLOYMENT, + api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`, + api_key: "fake-key", + input_cost_per_token: 0.01, + output_cost_per_token: 0.01, + }, + }, + }); + expect(res.ok(), `POST /model/new failed (${res.status()}): ${await res.text()}`).toBe(true); + + // /model/new returns once the row is written, but the router only picks the deployment up on its + // next refresh, so sending traffic straight away can still get "no healthy deployments". A ping + // that fails writes no spend log, so retrying it costs the ranking this test asserts nothing. + await expect + .poll( + async () => { + const ping = await request.post(`${rootPath()}/v1/chat/completions`, { + headers: { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }, + data: { model: modelName, messages: [{ role: "user", content: "readiness ping" }] }, + }); + return ping.ok(); + }, + { message: `deployment ${modelName} never became routable`, timeout: 60_000 }, + ) + .toBe(true); + + return { modelName }; +} + test.describe("Usage page", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); + const pricedDeployments: string[] = []; + + test.afterEach(async ({ request }) => { + // A deployment left behind keeps its custom pricing, so it goes on changing what later runs + // route and what they cost. Runs on the failure path too, which the test body would not. + // Resolved by name rather than by a returned id, so a create that persisted without answering + // 2xx is still cleaned up. /model/info serves the router, and /model/new answers 2xx even when + // its in-request router reload failed, so the search-backed listing is what covers a deployment + // that reached the database only. Absent from both means it never persisted. + const names = pricedDeployments.splice(0); + if (names.length === 0) return; + const auth = { Authorization: `Bearer ${masterKey()}`, "Content-Type": "application/json" }; + + type Lookup = + | { readonly listed: true; readonly id: string | undefined } + | { readonly listed: false; readonly status: number }; + + const idIn = async (path: string, name: string): Promise => { + const listed = await request.get(path, { headers: auth }); + if (!listed.ok()) return { listed: false, status: listed.status() }; + const deployments = ((await listed.json()).data ?? []) as { + model_name?: string; + model_info?: { id?: string }; + }[]; + return { listed: true, id: deployments.find((d) => d.model_name === name)?.model_info?.id }; + }; + + const remove = async (name: string, id: string) => { + const deleted = await request.post(`${rootPath()}/model/delete`, { headers: auth, data: { id } }); + expect(deleted.ok(), `POST /model/delete for ${name} (${deleted.status()})`).toBe(true); + }; + + for (const name of names) { + const fromRouter = await idIn(`${rootPath()}/model/info`, name); + if (fromRouter.listed && fromRouter.id !== undefined) { + await remove(name, fromRouter.id); + continue; + } + const search = encodeURIComponent(name); + const fromDb = await idIn(`${rootPath()}/v2/model/info?search=${search}`, name); + expect( + fromDb.listed, + `GET /v2/model/info?search=${search} (${fromDb.listed ? 200 : fromDb.status}), so ${name} could not be checked`, + ).toBe(true); + if (!fromDb.listed || fromDb.id === undefined) continue; + await remove(name, fromDb.id); + } + }); + test("Top Virtual Keys lists a key that served traffic, toggles views, and opens key info", async ({ page, request, @@ -39,8 +136,13 @@ test.describe("Usage page", () => { key_alias: alias, }); + // Top Virtual Keys ranks by spend, and every mock deployment costs $0, so once a run has more + // keys than the list shows, whether this one makes the cut is down to how ties happen to sort. + // Give it a priced deployment of its own so it earns its place. + const { modelName } = await createPricedDeployment(request, alias, pricedDeployments); + const requestId = await sendChatCompletion(request, { - model: CHAT_MODEL_A, + model: modelName, prompt: `usage ping for ${alias}`, apiKey: key, }); diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 05886e4b7f6..58cde4c8103 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -40,6 +40,24 @@ def prometheus_logger() -> PrometheusLogger: return PrometheusLogger() +@pytest.fixture +def known_model_router(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"}, + }, + { + "model_name": "us/azure/openai/gpt-5-mini", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "fake-key"}, + }, + ] + ) + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + yield router + + def create_standard_logging_payload() -> StandardLoggingPayload: return StandardLoggingPayload( id="test_id", @@ -741,7 +759,7 @@ async def test_async_log_failure_event(prometheus_logger): @pytest.mark.asyncio -async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger): +async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger, known_model_router): """LiteLLM-side reject (no deployment picked) routes the requested model into `requested_model` and skips the partial-outage flag.""" standard_logging_object = create_standard_logging_payload() @@ -786,7 +804,7 @@ async def test_async_log_failure_event_litellm_side_rate_limit(prometheus_logger @pytest.mark.asyncio -async def test_async_post_call_failure_hook(prometheus_logger): +async def test_async_post_call_failure_hook(prometheus_logger, known_model_router): """ Test for the async_post_call_failure_hook method @@ -1069,7 +1087,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): @pytest.mark.asyncio -async def test_log_success_fallback_event(prometheus_logger): +async def test_log_success_fallback_event(prometheus_logger, known_model_router): prometheus_logger.litellm_deployment_successful_fallbacks = MagicMock() original_model_group = "gpt-5-mini" @@ -1107,7 +1125,7 @@ async def test_log_success_fallback_event(prometheus_logger): @pytest.mark.asyncio -async def test_log_failure_fallback_event(prometheus_logger): +async def test_log_failure_fallback_event(prometheus_logger, known_model_router): prometheus_logger.litellm_deployment_failed_fallbacks = MagicMock() original_model_group = "gpt-5-mini" diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 66dbb29dba5..a86752c0172 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -841,7 +841,7 @@ def test_build_synthetic_response_events_covers_annotations_function_calls_and_r ) try: - events = streaming_module._build_synthetic_response_events( + events = streaming_module.build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=5, diff --git a/tests/llm_translation/conftest.py b/tests/llm_translation/conftest.py index 8532af2851c..567040c1d19 100644 --- a/tests/llm_translation/conftest.py +++ b/tests/llm_translation/conftest.py @@ -44,7 +44,11 @@ _VCR_AUTO_MARKER_SKIP_FILES = frozenset( {"test_vcr_redis_persister.py", "test_ws_vcr.py"} ) -_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = () +_VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ( + "test_nvidia_nim.py::test_embedding_nvidia_nim", + "test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[False]", + "test_litellm_proxy_provider.py::test_litellm_gateway_from_sdk_embedding[True]", +) _verbose_state = VerboseReporterState() diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index 485f78ed8c7..aa9f66f6665 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -163,6 +163,13 @@ _CAPS_NONE: FrozenSet[str] = frozenset() ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="claude-fable-5-1", + model="anthropic/claude-fable-5-1", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-fable-5", model="anthropic/claude-fable-5", @@ -222,6 +229,19 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="azure-claude-fable-5-1", + model="azure_ai/claude-fable-5-1", + mode="adaptive", + required_env=_AZURE_FOUNDRY_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5-1 has no deployment on the CI Microsoft Foundry " + "resource yet, so Foundry returns DeploymentNotFound and this cell " + "stays loud in CI. Remove this fail_reason once the deployment " + "exists." + ), + ), ModelEntry( alias="azure-claude-fable-5", model="azure_ai/claude-fable-5", @@ -268,6 +288,20 @@ AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="vertex-claude-fable-5-1", + model="vertex_ai/claude-fable-5-1", + mode="adaptive", + extra_params=(("vertex_location", "global"),), + required_env=_VERTEX_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5-1 availability on the CI Vertex project is not yet " + "confirmed for this brand-new release, so this cell stays loud in " + "CI until verified. Remove this fail_reason once the model is " + "confirmed available on the global Vertex endpoint." + ), + ), ModelEntry( alias="vertex-claude-fable-5", model="vertex_ai/claude-fable-5", @@ -332,6 +366,22 @@ VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="bedrock-claude-fable-5-1", + model="bedrock/converse/us.anthropic.claude-fable-5-1", + mode="adaptive", + extra_params=(("aws_region_name", "us-east-1"),), + required_env=_BEDROCK_REQ, + caps=_CAPS_XHIGH_MAX, + bedrock_effort_ceiling="xhigh", + unavailable_error="is not available for this account", + fail_reason=( + "claude-fable-5-1 access on the CI Bedrock account is not yet " + "confirmed for this brand-new release, so this cell stays loud in " + "CI until verified. Remove this fail_reason once the model is " + "enabled for the account." + ), + ), ModelEntry( alias="bedrock-claude-fable-5", model="bedrock/converse/us.anthropic.claude-fable-5", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index 517e3173b8c..714d544ecd6 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -201,8 +201,8 @@ async def test_reasoning_effort_grid( def test_grid_cell_count() -> None: - assert len(_PARAMS) == 31 * 11, ( - f"expected 341 cells (31 provider x model combos x 11 efforts), " + assert len(_PARAMS) == 35 * 11, ( + f"expected 385 cells (35 provider x model combos x 11 efforts), " f"got {len(_PARAMS)}" ) diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 1cb805bf9ba..8630259877d 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -5,6 +5,7 @@ from io import BytesIO from unittest.mock import AsyncMock +import httpx import litellm from litellm import completion, embedding import pytest @@ -92,44 +93,54 @@ async def test_litellm_gateway_from_sdk_embedding(is_async): litellm.set_verbose = True litellm._turn_on_debug() + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "my-vllm-model", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + if is_async: from openai import AsyncOpenAI - openai_client = AsyncOpenAI(api_key="fake-key") - mock_method = AsyncMock() - patch_target = openai_client.embeddings.create + openai_client = AsyncOpenAI( + api_key="fake-key", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + response = await litellm.aembedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) else: from openai import OpenAI - openai_client = OpenAI(api_key="fake-key") - mock_method = MagicMock() - patch_target = openai_client.embeddings.create + openai_client = OpenAI( + api_key="fake-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + response = litellm.embedding( + model="litellm_proxy/my-vllm-model", + input="Hello world", + client=openai_client, + api_base="my-custom-api-base", + ) - with patch.object(patch_target.__self__, patch_target.__name__, new=mock_method): - try: - if is_async: - await litellm.aembedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - else: - litellm.embedding( - model="litellm_proxy/my-vllm-model", - input="Hello world", - client=openai_client, - api_base="my-custom-api-base", - ) - except Exception as e: - print(e) + request_body = captured_bodies[0] + print("Request body - {}".format(request_body)) - mock_method.assert_called_once() - - print("Call KWARGS - {}".format(mock_method.call_args.kwargs)) - - assert "Hello world" == mock_method.call_args.kwargs["input"] - assert "my-vllm-model" == mock_method.call_args.kwargs["model"] + assert "Hello world" == request_body["input"] + assert "my-vllm-model" == request_body["model"] + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] @pytest.mark.parametrize("is_async", [False, True]) diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 7ee4f347f72..d5942e674d0 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -63,27 +63,39 @@ def test_embedding_nvidia_nim(): litellm.set_verbose = True from openai import OpenAI + captured_bodies = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "nvidia/nv-embedqa-e5-v5", + "usage": {"prompt_tokens": 6, "total_tokens": 6}, + }, + ) + client = OpenAI( api_key="fake-api-key", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), ) - with patch.object(client.embeddings.with_raw_response, "create") as mock_client: - try: - litellm.embedding( - model="nvidia_nim/nvidia/nv-embedqa-e5-v5", - input="What is the meaning of life?", - input_type="passage", - dimensions=1024, - client=client, - ) - except Exception as e: - print(e) - mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs - print("request_body: ", request_body) - assert request_body["input"] == "What is the meaning of life?" - assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" - assert request_body["extra_body"]["input_type"] == "passage" - assert request_body["dimensions"] == 1024 + response = litellm.embedding( + model="nvidia_nim/nvidia/nv-embedqa-e5-v5", + input="What is the meaning of life?", + input_type="passage", + dimensions=1024, + client=client, + ) + request_body = captured_bodies[0] + print("request_body: ", request_body) + assert request_body["input"] == "What is the meaning of life?" + assert request_body["model"] == "nvidia/nv-embedqa-e5-v5" + assert request_body["input_type"] == "passage" + assert request_body["dimensions"] == 1024 + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] def test_chat_completion_nvidia_nim_with_tools(): diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index a90a3df584e..7b03736920b 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1446,9 +1446,17 @@ def test_convert_to_anthropic_tool_invoke_sanitizes_invalid_ids(): def test_convert_to_anthropic_tool_invoke_server_tool(): """ - Test that server_tool_use (srvtoolu_) is reconstructed as server_tool_use. + Test that a server tool call (srvtoolu_) with no stored result is replayed + as a regular tool_use block. - Fixes: https://github.com/BerriAI/litellm/issues/17737 + A server_tool_use block is only valid when paired with its result block, so + an unpaired one must degrade to tool_use for Anthropic to accept the replay. + A paired call still becomes server_tool_use, covered by + test_convert_to_anthropic_tool_invoke_with_web_search_results. + + Context: https://github.com/BerriAI/litellm/issues/17737 (original + server_tool_use reconstruction) and LIT-6622 / PR #39144 (unpaired calls + degrade instead of 400ing at Anthropic). """ tool_calls = [ { @@ -1464,7 +1472,7 @@ def test_convert_to_anthropic_tool_invoke_server_tool(): result = convert_to_anthropic_tool_invoke(tool_calls) assert len(result) == 1 - assert result[0]["type"] == "server_tool_use" # NOT tool_use + assert result[0]["type"] == "tool_use" assert result[0]["id"] == "srvtoolu_01ABC123" assert result[0]["name"] == "web_search" assert result[0]["input"] == {"query": "elephant weight"} diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index ee93009a198..5535a62bb81 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -90,6 +90,7 @@ _VCR_INCOMPATIBLE_FILES = frozenset( # carry no real provider cost. _VCR_INCOMPATIBLE_NODEID_SUFFIXES: tuple[str, ...] = ( "test_router.py::test_router_text_completion_client", + "test_embedding.py::test_encoding_format_omitted_by_default_for_openai_sdk", ) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index aed2849f056..ee2ac14f498 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -3,6 +3,8 @@ import os import re import traceback +import httpx + import openai import pytest from dotenv import load_dotenv @@ -1255,56 +1257,42 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): assert sent_data["input"] == expected_payload_input -def test_encoding_format_defaults_to_float_for_openai_sdk(monkeypatch): +def test_encoding_format_omitted_by_default_for_openai_sdk(monkeypatch): """ - When encoding_format is not provided, LiteLLM sends `float` for OpenAI-path embeddings. + When encoding_format is not provided, LiteLLM leaves it out of the upstream request. Optional global override: `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT`. """ monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - # Create a mock client instance - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance + captured_bodies = [] - # Mock the embeddings.with_raw_response.create method - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", + def handler(request: httpx.Request) -> httpx.Response: + captured_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-ada-002", "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} - - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response + }, ) - # Call the embedding function without encoding_format - response = embedding( - model="text-embedding-ada-002", - input="Hello world", - ) + client = openai.OpenAI( + api_key="sk-test", http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) - # Get the call arguments to verify what was sent to OpenAI SDK - call_args = mock_client_instance.embeddings.with_raw_response.create.call_args - assert ( - call_args is not None - ), "OpenAI SDK embeddings.create should have been called" + response = embedding( + model="text-embedding-ada-002", + input="Hello world", + api_key="sk-test", + client=client, + ) - call_kwargs = call_args[1] # Get kwargs - - assert "encoding_format" in call_kwargs - assert ( - call_kwargs["encoding_format"] == "float" - ), "encoding_format should default to float when not provided by user" - - print("✅ PASS: encoding_format='float' is correctly passed to OpenAI SDK") + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert "encoding_format" not in captured_bodies[0], ( + "encoding_format should be omitted from the upstream request when not provided by user" + ) def test_encoding_format_explicit_value_preserved(): diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 8370046446d..e6392cda406 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -5,7 +5,7 @@ import traceback from typing import Any import httpx -from openai import AsyncOpenAI, AuthenticationError, BadRequestError, OpenAIError, RateLimitError +from openai import AsyncAzureOpenAI, AsyncOpenAI, AuthenticationError, AzureOpenAI, BadRequestError, OpenAIError, RateLimitError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -895,7 +895,12 @@ def _pre_call_utils( ): if call_type == "embedding": data["input"] = "Hello world!" - mapped_target: Any = client.embeddings.with_raw_response + if isinstance(client, (AzureOpenAI, AsyncAzureOpenAI)): + mapped_target: Any = client.embeddings.with_raw_response + patched_attr = "create" + else: + mapped_target = client + patched_attr = "post" if sync_mode: original_function = litellm.embedding else: @@ -905,6 +910,7 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.chat.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.completion else: @@ -914,12 +920,13 @@ def _pre_call_utils( if streaming is True: data["stream"] = True mapped_target = client.completions.with_raw_response # type: ignore + patched_attr = "create" if sync_mode: original_function = litellm.text_completion else: original_function = litellm.atext_completion - return data, original_function, mapped_target + return data, original_function, mapped_target, patched_attr def _pre_call_utils_httpx( @@ -1003,7 +1010,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str ) data = {"model": model} - data, original_function, mapped_target = _pre_call_utils( + data, original_function, mapped_target, patched_attr = _pre_call_utils( call_type=call_type, data=data, client=openai_client, @@ -1049,7 +1056,7 @@ async def test_exception_with_headers(sync_mode, provider, model, call_type, str with patch.object( mapped_target, - "create", + patched_attr, side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index cc6209f2bf9..ebad0fbafc5 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -155,6 +155,11 @@ def test_default_api_base(): continue elif provider == "github" and other_provider.value == "azure": continue + elif ( + provider in ("qwencloud", "qwen_ai_platform") + and other_provider.value == "dashscope" + ): + continue assert other_provider.value not in api_base.replace("/openai", "") diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 370c43f8f44..c714bb4f9a7 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -2032,8 +2032,8 @@ def test_router_dynamic_cooldown_correct_retry_after_time(): raise exception with patch.object( - openai_client.embeddings.with_raw_response, - "create", + openai_client, + "post", side_effect=_return_exception, ): new_retry_after_mock_client = MagicMock(return_value=-1) diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index aadaadd510e..6438525706a 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport -from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult +from mcp.types import CallToolResult as MCPCallToolResult +from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Tool as MCPTool def test_mcp_client_uses_configurable_default_timeout(): @@ -185,6 +187,80 @@ class TestMCPClientUnitTests: mock_session_instance.initialize.assert_called_once() mock_session_instance.list_tools.assert_called_once() + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + @patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + async def test_list_tools_follows_next_cursor_until_exhausted( + self, + mock_session_class, + mock_transport, + ): + """Test listing tools follows MCP pagination cursors until exhausted.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + first_page_tools = [ + MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100) + ] + second_page_tool = MCPTool( + name="tool_100", + description="Tool 100", + inputSchema={}, + ) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult(tools=first_page_tools, nextCursor="page-2"), + ListToolsResult(tools=[second_page_tool]), + ] + + client = MCPClient("http://example.com") + result = await client.list_tools() + + assert result == [*first_page_tools, second_page_tool] + assert mock_session_instance.list_tools.call_count == 2 + second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + @patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + async def test_list_tools_swallows_mid_walk_error_without_raise_on_error( + self, + mock_session_class, + mock_transport, + ): + """Test a mid-walk failure returns [] when raise_on_error is False.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})], + nextCursor="page-2", + ), + RuntimeError("transient upstream failure"), + ] + + client = MCPClient("http://example.com") + result = await client.list_tools() + + assert result == [] + assert mock_session_instance.list_tools.call_count == 2 + @pytest.mark.asyncio @patch.object(mcp_client_module, "streamable_http_client") @patch.object(mcp_client_module, "ClientSession") diff --git a/tests/ocr_tests/test_ocr_vertex_ai.py b/tests/ocr_tests/test_ocr_vertex_ai.py index 1ba5b9d0883..1842eb063a5 100644 --- a/tests/ocr_tests/test_ocr_vertex_ai.py +++ b/tests/ocr_tests/test_ocr_vertex_ai.py @@ -5,9 +5,11 @@ Note: Vertex AI OCR automatically converts URLs to base64 data URIs since the Vertex AI endpoint doesn't have internet access. """ -import os import json +import os import tempfile +from typing import Final + import pytest from base_ocr_unit_tests import BaseOCRTest @@ -139,3 +141,19 @@ def test_vertex_ai_ocr_routing(): assert isinstance( deepseek_variant, VertexAIDeepSeekOCRConfig ), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig" + + +@pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas")) +def test_deepseek_request_uses_single_provider_namespace(model: str) -> None: + from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( + VertexAIDeepSeekOCRConfig, + ) + + request: Final = VertexAIDeepSeekOCRConfig().transform_ocr_request( + model=model, + document={"type": "image_url", "image_url": "data:image/png;base64,AA=="}, + optional_params={}, + headers={}, + ) + + assert request.data["model"] == "deepseek-ai/deepseek-ocr-maas" diff --git a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py index 091ea106b91..fd95b7fa8f2 100644 --- a/tests/pass_through_unit_tests/test_websearch_interception_e2e.py +++ b/tests/pass_through_unit_tests/test_websearch_interception_e2e.py @@ -937,6 +937,14 @@ async def test_pre_request_hook_modifies_request_body(): print("✅ WebSearchInterceptionLogger initialized") + mock_router = MagicMock() + mock_router.search_tools = [ + { + "search_tool_name": "test-search-tool", + "litellm_params": {"search_provider": "tavily"}, + } + ] + # Track what actually gets sent to the API captured_request = {} @@ -987,6 +995,9 @@ async def test_pre_request_hook_modifies_request_body(): with patch( "litellm.llms.anthropic.experimental_pass_through.messages.handler.anthropic_messages_handler", side_effect=mock_anthropic_messages_handler, + ), patch( # test-quality-ok: the hook imports this process-global router at call time; no injection seam exists to register search_tools + "litellm.proxy.proxy_server.llm_router", + mock_router, ): print( diff --git a/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py b/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py new file mode 100644 index 00000000000..ed21734c5fc --- /dev/null +++ b/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py @@ -0,0 +1,58 @@ +"""Image-level check that the built proxy image can import the Bedrock realtime SDK. + +Bedrock Nova Sonic (`/v1/realtime`) imports `aws_sdk_bedrock_runtime` lazily on the +first session, so an image whose `uv sync` stages skip the `bedrock-realtime` extra +boots, passes health checks, and then fails every Nova Sonic session with +"Missing aws_sdk_bedrock_runtime". Importing inside the built image is what catches +that class of regression (missing extra, lockfile drift, a stage that syncs a +different set of extras), which a static Dockerfile check cannot. + +Gated on LITELLM_IMAGE like the other image checks in this directory; exercised +where an image has been built (the image-scan workflow). Requires a working docker CLI. +""" + +import os +import shutil +import subprocess +from typing import Final + +import pytest + +IMAGE: Final = os.getenv("LITELLM_IMAGE") +NON_ROOT_UID: Final = "12345:0" +IMPORT_PROBE: Final = "import aws_sdk_bedrock_runtime, smithy_aws_core; print('bedrock-realtime ok')" + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def test_image_imports_bedrock_realtime_sdk(): + assert IMAGE is not None + + probe: Final = subprocess.run( + [ + "docker", + "run", + "--rm", + "--network", + "none", + "--user", + NON_ROOT_UID, + "--entrypoint", + "python", + IMAGE, + "-c", + IMPORT_PROBE, + ], + capture_output=True, + text=True, + check=False, + ) + + assert probe.returncode == 0 and "bedrock-realtime ok" in probe.stdout, ( + f"{IMAGE} cannot import aws_sdk_bedrock_runtime as uid {NON_ROOT_UID}, so Bedrock Nova Sonic " + "/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'. Is `--extra bedrock-realtime` " + f"passed to every `uv sync` in its Dockerfile?\nstdout:\n{probe.stdout}\nstderr:\n{probe.stderr}" + ) diff --git a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py new file mode 100644 index 00000000000..0c4d1dfc21e --- /dev/null +++ b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py @@ -0,0 +1,402 @@ +""" +Unit tests for safeguard-refusal fallback on the /v1/messages router surface. + +An Anthropic safeguard refusal is an HTTP 200 whose body carries +stop_reason "refusal" plus a stop_details object; the router converts it +into a ContentPolicyViolationError so the content-policy fallback chain +runs, but only when a matching fallback is configured. A plain refusal +without stop_details, or any refusal with nothing configured, must reach +the client byte-identical. + +The upstream is faked at the HTTP boundary by intercepting the third-party +transport (httpx.AsyncClient.send), so requests run litellm's real +transformation, allowlist, and streaming pipeline end to end. +""" + +import json +from typing import Any, AsyncIterator +from unittest.mock import patch + +import httpx +import pytest + +from litellm import Router +from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + record_pre_routing_selection, +) + +REFUSAL_RESPONSE: dict[str, Any] = { + "id": "msg_refusal", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"category": "cyber", "explanation": "flagged"}, + "usage": {"input_tokens": 25, "output_tokens": 1}, +} + +PLAIN_REFUSAL_RESPONSE: dict[str, Any] = {k: v for k, v in REFUSAL_RESPONSE.items() if k != "stop_details"} + +OK_RESPONSE: dict[str, Any] = { + "id": "msg_ok", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 25, "output_tokens": 2}, +} + + +def _sse(event: str, data: dict[str, Any]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "refusal", "stop_details": {"category": "cyber"}}, + "usage": {"output_tokens": 1}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), +) + +OK_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**OK_RESPONSE, "stop_reason": None}}), + _sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + _sse("message_stop", {"type": "message_stop"}), +) + + +def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]: + """Split each frame's data line in half, modeling a transport chunk boundary.""" + return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :])) + + +class _FrameStream(httpx.AsyncByteStream): + def __init__(self, frames: tuple[bytes, ...]) -> None: + self._frames = frames + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + yield frame + + async def aclose(self) -> None: + return None + + +class FakeAnthropicUpstream: + """Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable + models, answers on others. The router deliberately does not forward caller-injected + clients, so the transport is the seam that exercises the real litellm pipeline.""" + + def __init__( + self, + refusal_body: dict[str, Any] = REFUSAL_RESPONSE, + refusal_frames: tuple[bytes, ...] = REFUSAL_STREAM_FRAMES, + ) -> None: + self.refusal_body = refusal_body + self.refusal_frames = refusal_frames + self.calls: list[str] = [] + self.bodies: list[dict[str, Any]] = [] + + async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response: + body = json.loads(request.content or b"{}") + model = body.get("model", "") + self.calls.append(model) + self.bodies.append(body) + refuses = "fable" in model + if body.get("stream"): + frames = self.refusal_frames if refuses else OK_STREAM_FRAMES + return httpx.Response( + 200, + stream=_FrameStream(frames), + headers={"content-type": "text/event-stream"}, + request=request, + ) + return httpx.Response(200, json=self.refusal_body if refuses else OK_RESPONSE, request=request) + + def install(self): + async def _send(_client: httpx.AsyncClient, request: httpx.Request, **kwargs: Any) -> httpx.Response: + return await self.send(request, **kwargs) + + return patch("httpx.AsyncClient.send", new=_send) + + +FABLE_TIER = { + "model_name": "fable-tier", + "litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"}, +} +OPUS_TARGET = { + "model_name": "opus-target", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"}, +} + + +def _router(content_policy_fallbacks: list | None) -> Router: + return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks) + + +async def _collect(stream: AsyncIterator[bytes]) -> bytes: + return b"".join([chunk async for chunk in stream]) + + +@pytest.mark.asyncio +async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "end_turn" + assert response["id"] == "msg_ok" + assert len(fake.calls) == 2 + assert "claude-opus-5" in fake.calls[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content_policy_fallbacks, upstream_body", + [ + (None, REFUSAL_RESPONSE), + ([{"unrelated-group": ["opus-target"]}], REFUSAL_RESPONSE), + ([{"fable-tier": ["opus-target"]}], PLAIN_REFUSAL_RESPONSE), + ], + ids=["nothing-configured", "row-for-other-group", "refusal-without-stop-details"], +) +async def test_non_streaming_refusal_passes_through_untouched(content_policy_fallbacks, upstream_body): + fake = FakeAnthropicUpstream(refusal_body=upstream_body) + router = _router(content_policy_fallbacks=content_policy_fallbacks) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "refusal" + assert response.get("stop_details") == upstream_body.get("stop_details") + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_with_fallback_row_streams_fallback_frames(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_split_across_chunks_still_falls_back(): + fake = FakeAnthropicUpstream(refusal_frames=_split_frames_mid_data_line(REFUSAL_STREAM_FRAMES)) + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_without_fallback_row_passes_frames_through(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=None) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert b"stop_details" in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): + """The pre-routing hook's tier stamp must reach the mid-stream fallback lookup even when the + request carries no metadata bucket at all (the snapshot is taken before the request runs).""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + stream = await router.aanthropic_messages( + model="smart-router", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"forged-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "hi"}], + litellm_metadata={PRE_ROUTING_SELECTED_MODEL_KEY: "forged-tier"}, + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_tier_stamp_never_reaches_provider_bound_metadata(): + """On /v1/messages the top-level metadata dict is Anthropic's own request field, so the + routed-tier stamp must never appear in any upstream body even when the client sends one.""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + response = await router.aanthropic_messages( + model="smart-router", + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + metadata={"user_id": "u1"}, + ) + + assert response["stop_reason"] == "end_turn" + assert len(fake.bodies) == 2 + for body in fake.bodies: + assert body.get("metadata") == {"user_id": "u1"} + + +def test_record_pre_routing_selection_writes_only_the_internal_bucket(): + """The Anthropic request's own metadata field must never carry the tier stamp.""" + kwargs = {"metadata": {"user_id": "u1"}, "litellm_metadata": {}} + + record_pre_routing_selection(kwargs, "tier-x") + + assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"} + assert kwargs["metadata"] == {"user_id": "u1"} + + +def test_refusal_gate_keys_on_pre_routing_tier_stamp(): + router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) + + def anthropic_messages(**kwargs: Any) -> None: + return None + + refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs=refusal_kwargs, + ) + is True + ) + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) + + +def test_has_content_policy_fallback_default_fallbacks_arm(): + router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}]) + + assert router._has_content_policy_fallback("any-group", {}) is True + assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + + +def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested(): + router = _router(content_policy_fallbacks=None) + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}] + + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier1", "smart-router") + ) == ["backup-a"] + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier9", "smart-router") + ) == ["backup-b"] + assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None + + +def test_refusal_gate_ignores_other_generic_call_types(): + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + def aresponses(**kwargs: Any) -> None: + return None + + assert ( + router._should_raise_anthropic_refusal_error( + model="fable-tier", + original_generic_function=aresponses, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index dcd2e9edf7b..7dbac243d55 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -2294,6 +2294,7 @@ def search_tools(): "search_provider": "perplexity", "api_key": "test-api-key", "api_base": "https://api.perplexity.ai", + "mode": "turbo", }, }, { @@ -2302,6 +2303,7 @@ def search_tools(): "search_provider": "perplexity", "api_key": "test-api-key-2", "api_base": "https://api.perplexity.ai", + "mode": "turbo", }, }, ] @@ -2393,6 +2395,7 @@ async def test_asearch_with_fallbacks_helper(search_tools): assert "search_provider" in kwargs assert kwargs["search_provider"] == "perplexity" assert "api_key" in kwargs + assert kwargs["mode"] == "turbo" assert kwargs["query"] == "helper test query" return mock_response diff --git a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py index c0c720bbaf6..953f028af3c 100644 --- a/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py +++ b/tests/test_litellm/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -1,3 +1,4 @@ +import base64 from typing import Final from unittest.mock import MagicMock @@ -8,8 +9,17 @@ from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ( SpeechToCompletionBridgeTransformationHandler, ) +from litellm.types.utils import ChatCompletionAudioResponse, Choices, Message, ModelResponse GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" +PCM_BYTES: Final = b"\x01\x02\x03\x04" * 6 + + +def _model_response(model: str, pcm: bytes) -> ModelResponse: + audio: Final = ChatCompletionAudioResponse( + data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello" + ) + return ModelResponse(model=model, choices=[Choices(message=Message(content=None, audio=audio))]) def _bridge_request(response_format: str | None) -> dict: @@ -28,7 +38,7 @@ def _bridge_request(response_format: str | None) -> dict: ) -@pytest.mark.parametrize("response_format", ["wav", "mp3", "pcm", None]) +@pytest.mark.parametrize("response_format", ["wav", "pcm", None]) def test_gemini_tts_request_keeps_speech_response_format_out_of_chat_params(response_format: str | None) -> None: request: Final = _bridge_request(response_format) @@ -60,3 +70,48 @@ def test_non_gemini_request_forwards_speech_response_format_as_audio_format() -> assert "response_format" not in request assert request["audio"] == {"voice": "alloy", "format": "wav"} + + +@pytest.mark.parametrize("response_format", ["mp3", "flac", "opus", "aac"]) +def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_format: str) -> None: + with pytest.raises(litellm.BadRequestError) as excinfo: + _bridge_request(response_format) + + assert excinfo.value.status_code == 400 + assert response_format in str(excinfo.value) + assert "pcm" in str(excinfo.value) + assert "wav" in str(excinfo.value) + + +def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + response_format="pcm", + ) + + assert response.response.content == PCM_BYTES + assert response.response.headers["content-type"] == "audio/pcm" + + +@pytest.mark.parametrize("response_format", ["wav", None]) +def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: str | None) -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + response_format=response_format, + ) + + body: Final = response.response.content + assert body[:4] == b"RIFF" + assert body[8:12] == b"WAVE" + assert body[44:] == PCM_BYTES + assert response.response.headers["content-type"] == "audio/wav" + + +def test_non_gemini_response_keeps_original_bytes_and_mpeg_content_type() -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response("gpt-4o-audio-preview", PCM_BYTES), + response_format="mp3", + ) + + assert response.response.content == PCM_BYTES + assert response.response.headers["content-type"] == "audio/mpeg" diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/test_litellm/experimental_mcp_client/test_tools.py index 89f67452f29..6645b06664d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_tools.py +++ b/tests/test_litellm/experimental_mcp_client/test_tools.py @@ -8,11 +8,13 @@ from mcp.types import ( CallToolRequestParams, CallToolResult, ListToolsResult, + PaginatedRequestParams, TextContent, ) from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.tools import ( + list_tools_with_pagination, transform_mcp_tool_to_anthropic_tool, _get_function_arguments, _normalize_mcp_input_schema, @@ -106,6 +108,134 @@ async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result mock_session.list_tools.assert_called_once() +@pytest.mark.asyncio() +async def test_load_mcp_tools_follows_pagination(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[ + MCPTool(name="tool_a", description="a", inputSchema={}), + MCPTool(name="tool_b", description="b", inputSchema={}), + ], + nextCursor="page-2", + ), + ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]), + ] + result = await load_mcp_tools(mock_session, format="mcp") + assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"] + assert mock_session.list_tools.call_count == 2 + second_call_params = mock_session.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch): + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2) + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="page-2", + ), + ListToolsResult( + tools=[MCPTool(name="tool_1", description="1", inputSchema={})], + nextCursor="page-3", + ), + ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + assert mock_session.list_tools.call_count == 2 + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_on_repeated_cursor(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="same-cursor", + ), + ListToolsResult( + tools=[MCPTool(name="tool_1", description="1", inputSchema={})], + nextCursor="same-cursor", + ), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + assert mock_session.list_tools.call_count == 2 + + +@pytest.mark.asyncio() +async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="", + ), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0"] + mock_session.list_tools.assert_called_once() + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch): + import anyio + + from litellm.experimental_mcp_client.tools import list_tools_with_pagination + + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2) + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.2) + + async def slow_page(params=None): + await anyio.sleep(0.15) + idx = int(params.cursor) if params is not None else 0 + return ListToolsResult( + tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})], + nextCursor=str(idx + 1), + ) + + mock_session.list_tools = slow_page + result = await list_tools_with_pagination(mock_session) + + assert [tool.name for tool in result] == ["tool_0"] + + +@pytest.mark.asyncio() +async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_session, monkeypatch): + import anyio + + from litellm.experimental_mcp_client.tools import list_tools_with_pagination + + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1) + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.1) + + async def slow_page(params=None): + await anyio.sleep(0.15) + idx = int(params.cursor) if params is not None else 0 + tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})] + if idx == 0: + return ListToolsResult(tools=tools, nextCursor="1") + return ListToolsResult(tools=tools) + + mock_session.list_tools = slow_page + result = await list_tools_with_pagination(mock_session, listing_deadline=2.0) + + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + + +@pytest.mark.asyncio() +async def test_load_mcp_tools_openai_format_spans_pages(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_a", description="a", inputSchema={})], + nextCursor="page-2", + ), + ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]), + ] + result = await load_mcp_tools(mock_session, format="openai") + assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"] + + def test_get_function_arguments(): # Test with string arguments function = {"arguments": '{"test": "value"}'} diff --git a/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py b/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py new file mode 100644 index 00000000000..45e1acecec8 --- /dev/null +++ b/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py @@ -0,0 +1,193 @@ +import datetime +from typing import Final +from unittest.mock import AsyncMock, patch + +import pytest +from pydantic import ValidationError + +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.SlackAlerting.user_spend_alerts import ( + UserSpendRow, + evaluate_user_spend, +) +from litellm.types.integrations.slack_alerting import ( + DEFAULT_ALERT_TYPES, + AlertType, + SlackAlertingArgs, +) + +TODAY: Final = datetime.date(2026, 8, 15) + + +def _row( + daily_spend: float = 0.0, + monthly_spend: float = 0.0, + baseline_spend: float = 0.0, +) -> UserSpendRow: + return UserSpendRow( + user_id="user-1", + daily_spend=daily_spend, + monthly_spend=monthly_spend, + baseline_spend=baseline_spend, + ) + + +def _evaluate(row: UserSpendRow, args: SlackAlertingArgs, thresholds: bool = True, anomalies: bool = True): + return evaluate_user_spend( + row=row, + args=args, + today=TODAY, + thresholds_enabled=thresholds, + anomalies_enabled=anomalies, + ) + + +def test_daily_threshold_crossed(): + args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0) + events: Final = _evaluate(_row(daily_spend=75.0, monthly_spend=75.0), args) + assert [e.kind for e in events] == ["daily_threshold"] + assert "`$75.00`" in events[0].message + assert "`$50.00`" in events[0].message + assert events[0].alert_type == AlertType.user_spend_thresholds + assert events[0].cache_key == "user_spend_alert_daily_user-1_2026-08-15" + + +def test_daily_threshold_not_crossed(): + args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0) + assert _evaluate(_row(daily_spend=49.99, monthly_spend=49.99), args) == () + + +def test_thresholds_unset_by_default(): + args: Final = SlackAlertingArgs(spend_anomaly_min_spend=1000.0) + assert _evaluate(_row(daily_spend=999.0, monthly_spend=999.0), args) == () + + +def test_monthly_threshold_crossed(): + args: Final = SlackAlertingArgs(monthly_spend_per_user_threshold=200.0, spend_anomaly_min_spend=1000.0) + events: Final = _evaluate(_row(daily_spend=5.0, monthly_spend=250.0), args) + assert [e.kind for e in events] == ["monthly_threshold"] + assert events[0].cache_key == "user_spend_alert_monthly_user-1_2026-08" + + +def test_thresholds_disabled_suppresses_threshold_events(): + args: Final = SlackAlertingArgs( + daily_spend_per_user_threshold=50.0, + monthly_spend_per_user_threshold=200.0, + spend_anomaly_min_spend=1000.0, + ) + assert _evaluate(_row(daily_spend=75.0, monthly_spend=250.0), args, thresholds=False) == () + + +def test_anomaly_detected_above_multiple_of_baseline(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + events: Final = _evaluate( + _row(daily_spend=70.0, monthly_spend=100.0, baseline_spend=70.0), args + ) + assert [e.kind for e in events] == ["anomaly"] + assert events[0].alert_type == AlertType.user_spend_anomalies + assert "`$10.00`" in events[0].message + assert events[0].cache_key == "user_spend_alert_anomaly_user-1_2026-08-15" + + +def test_no_anomaly_within_baseline_multiple(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert ( + _evaluate(_row(daily_spend=25.0, monthly_spend=100.0, baseline_spend=70.0), args) == () + ) + + +def test_no_anomaly_below_min_spend_floor(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert _evaluate(_row(daily_spend=9.0, monthly_spend=9.0, baseline_spend=0.1), args) == () + + +def test_anomaly_for_new_user_without_baseline(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + events: Final = _evaluate(_row(daily_spend=15.0, monthly_spend=15.0), args) + assert [e.kind for e in events] == ["anomaly"] + + +def test_sparse_baseline_averages_over_full_window(): + args: Final = SlackAlertingArgs( + spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0, spend_anomaly_baseline_days=7 + ) + events: Final = _evaluate(_row(daily_spend=13.0, monthly_spend=20.0, baseline_spend=7.0), args) + assert [e.kind for e in events] == ["anomaly"] + + +def test_anomalies_not_in_default_alert_types(): + assert AlertType.user_spend_anomalies not in DEFAULT_ALERT_TYPES + assert AlertType.user_spend_thresholds in DEFAULT_ALERT_TYPES + + +def test_invalid_config_rejected(): + with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"): + SlackAlertingArgs(daily_spend_per_user_threshold=0) + with pytest.raises(ValidationError, match="spend_anomaly_baseline_days"): + SlackAlertingArgs(spend_anomaly_baseline_days=0) + with pytest.raises(ValidationError, match="user_spend_check_interval"): + SlackAlertingArgs(user_spend_check_interval=10) + + +def test_non_finite_config_rejected(): + with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"): + SlackAlertingArgs(daily_spend_per_user_threshold=float("inf")) + with pytest.raises(ValidationError, match="spend_anomaly_multiplier"): + SlackAlertingArgs(spend_anomaly_multiplier=float("nan")) + with pytest.raises(ValidationError, match="spend_anomaly_min_spend"): + SlackAlertingArgs(spend_anomaly_min_spend=float("inf")) + with pytest.raises(ValidationError, match="user_spend_check_interval"): + SlackAlertingArgs(user_spend_check_interval=float("inf")) + + +def test_anomalies_disabled_suppresses_anomaly_events(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert _evaluate(_row(daily_spend=500.0, monthly_spend=500.0), args, anomalies=False) == () + + +@pytest.mark.asyncio +async def test_send_user_spend_alerts_sends_and_dedupes(): + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alerting_args={"daily_spend_per_user_threshold": 50.0, "spend_anomaly_min_spend": 1000.0}, + ) + mock_prisma: Final = AsyncMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "user_id": "user-1", + "daily_spend": 75.0, + "monthly_spend": 75.0, + "baseline_spend": 0.0, + }, + { + "user_id": "user-2", + "daily_spend": 60.0, + "monthly_spend": 60.0, + "baseline_spend": 0.0, + }, + ] + ) + with patch.object(slack_alerting, "send_alert", new_callable=AsyncMock) as mock_send_alert: + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + assert mock_send_alert.call_count == 1 + sent_kwargs: Final = mock_send_alert.call_args.kwargs + assert sent_kwargs["alert_type"] == AlertType.user_spend_thresholds + assert "User Daily Spend Threshold Crossed" in sent_kwargs["message"] + assert "`user-1`" in sent_kwargs["message"] + assert "`user-2`" in sent_kwargs["message"] + + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + assert mock_send_alert.call_count == 1 + + +@pytest.mark.asyncio +async def test_send_user_spend_alerts_noop_when_alert_types_disabled(): + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alert_types=[AlertType.budget_alerts], + alerting_args={"daily_spend_per_user_threshold": 50.0}, + ) + mock_prisma: Final = AsyncMock() + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + mock_prisma.db.query_raw.assert_not_called() diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py new file mode 100644 index 00000000000..2d0605e3b7f --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -0,0 +1,469 @@ +""" +Regression tests for the Datadog LLM Observability payload schema (issue #35786). + +Datadog renders tool calls, tool results and prompt-cache savings only from the fields its +own schema names. These assert on the payload `create_llm_obs_payload` actually hands the +intake, so a regression that moves data back into `meta.metadata` fails here. + +Fixtures mirror what a live proxy run recorded on the callback, including the provider +spelling of prompt-cache counts (`prompt_tokens_details.cached_tokens`). +""" + +import json +import os +from datetime import datetime, timedelta +from typing import Any +from unittest.mock import patch + +import pytest + +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +TOOL_DEFINITION: dict[str, Any] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + +ASSISTANT_TOOL_CALL: dict[str, Any] = { + "id": "call_abc123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris","unit":"c"}'}, +} + + +@pytest.fixture +def logger() -> DataDogLLMObsLogger: + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + return DataDogLLMObsLogger() + + +NOT_GIVEN: Any = object() + + +def build_payload( + messages: Any = NOT_GIVEN, + response_message: dict[str, Any] | None = None, + usage_object: dict[str, Any] | None = None, + model_parameters: dict[str, Any] | None = None, + prompt_tokens: int = 4447, +) -> dict[str, Any]: + return { + "standard_logging_object": { + "call_type": "acompletion", + "messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages, + "response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]}, + "model_parameters": model_parameters or {}, + "metadata": {"usage_object": usage_object} if usage_object is not None else {}, + "prompt_tokens": prompt_tokens, + "completion_tokens": 507, + "total_tokens": prompt_tokens + 507, + "response_cost": 0.02, + "status": "success", + }, + "litellm_params": {"metadata": {}}, + } + + +def build(logger: DataDogLLMObsLogger, **kwargs: Any) -> dict[str, Any]: + """Build a span and read it back as the JSON the intake receives, not as Python objects.""" + start = datetime(2026, 9, 1, 12, 0, 0) + payload = logger.create_llm_obs_payload(build_payload(**kwargs), start, start + timedelta(seconds=2)) + return json.loads(safe_dumps(payload)) + + +def test_output_tool_calls_use_the_datadog_tool_call_schema(logger: DataDogLLMObsLogger) -> None: + """Datadog reads name/arguments/tool_id off the tool call; OpenAI nests them under `function`.""" + payload = build( + logger, + response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + + message = payload["meta"]["output"]["messages"][0] + assert message["tool_calls"] == [ + { + "name": "get_weather", + "arguments": {"city": "Paris", "unit": "c"}, + "tool_id": "call_abc123", + "type": "function", + } + ] + assert "function" not in message["tool_calls"][0] + + +def test_tool_calls_are_not_duplicated_into_metadata(logger: DataDogLLMObsLogger) -> None: + """The flat `output_tool_calls.*` keys were a second copy of a fact that now has its own field.""" + payload = build( + logger, + response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + + assert [key for key in payload["meta"]["metadata"] if "tool_calls." in key] == [] + + +def test_tool_result_message_links_back_to_its_tool_call(logger: DataDogLLMObsLogger) -> None: + """Datadog pairs a result with its call through tool_id, and names the tool from the call.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + {"role": "tool", "tool_call_id": "call_abc123", "content": '{"temp_c": 18}'}, + ], + ) + + tool_message = payload["meta"]["input"]["messages"][2] + assert tool_message["tool_results"] == [ + {"name": "get_weather", "result": '{"temp_c": 18}', "tool_id": "call_abc123", "type": "function"} + ] + + +def test_tool_result_without_a_matching_call_still_reports_its_id(logger: DataDogLLMObsLogger) -> None: + """A truncated conversation loses the call, so the name is unknown but the link must survive.""" + payload = build( + logger, + messages=[{"role": "tool", "tool_call_id": "call_orphan", "content": "42"}], + ) + + assert payload["meta"]["input"]["messages"][0]["tool_results"] == [ + {"name": "", "result": "42", "tool_id": "call_orphan", "type": "function"} + ] + + +def test_cache_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None: + """ + Datadog charts cache savings from span metrics; nested usage_object is not read for it. + + litellm's normalized prompt count includes both cache categories, so the three cache + metrics must partition input_tokens: read + write + non_cached == input. + """ + payload = build( + logger, + usage_object={"prompt_tokens_details": {"cached_tokens": 4300, "cache_write_tokens": 95}}, + ) + + metrics = payload["metrics"] + assert metrics["cache_read_input_tokens"] == 4300.0 + assert metrics["cache_write_input_tokens"] == 95.0 + assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0 + assert ( + metrics["cache_read_input_tokens"] + metrics["cache_write_input_tokens"] + metrics["non_cached_input_tokens"] + == metrics["input_tokens"] + ) + + +def test_cache_write_tokens_are_not_counted_as_non_cached(logger: DataDogLLMObsLogger) -> None: + """A cache-priming request must not report its primed prefix as full-price uncached input.""" + payload = build(logger, usage_object={"prompt_tokens_details": {"cache_write_tokens": 4000}}) + + assert payload["metrics"]["cache_write_input_tokens"] == 4000.0 + assert payload["metrics"]["non_cached_input_tokens"] == 4447.0 - 4000.0 + assert "cache_read_input_tokens" not in payload["metrics"] + + +def test_a_fully_cached_request_reports_a_zero_non_cached_count(logger: DataDogLLMObsLogger) -> None: + """Zero residual is real data: everything was served from cache. Inconsistent counts clamp to it.""" + payload = build( + logger, + usage_object={"prompt_tokens_details": {"cached_tokens": 4352, "cache_write_tokens": 95}}, + ) + + assert payload["metrics"]["non_cached_input_tokens"] == 0.0 + + +def test_anthropic_top_level_cache_keys_are_read(logger: DataDogLLMObsLogger) -> None: + """A raw Anthropic usage dict records the counts top level, not under prompt_tokens_details.""" + payload = build( + logger, + usage_object={"cache_read_input_tokens": 4300, "cache_creation_input_tokens": 95}, + ) + + metrics = payload["metrics"] + assert metrics["cache_read_input_tokens"] == 4300.0 + assert metrics["cache_write_input_tokens"] == 95.0 + assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0 + + +def test_cache_metrics_come_from_the_normalized_field_not_the_anthropic_one(logger: DataDogLLMObsLogger) -> None: + """ + litellm normalizes every provider's cache counters into prompt_tokens_details. + + A real cached request from a non-Anthropic provider carries only `cached_tokens`, so + reading the Anthropic-specific `cache_read_input_tokens` key reports nothing for it. + """ + payload = build( + logger, + usage_object={"prompt_tokens_details": {"audio_tokens": None, "cached_tokens": 4096}}, + prompt_tokens=4335, + ) + + assert payload["metrics"]["cache_read_input_tokens"] == 4096.0 + assert payload["metrics"]["non_cached_input_tokens"] == 4335.0 - 4096.0 + + +@pytest.mark.parametrize( + "usage_object", + [ + {"prompt_tokens_details": {"cache_write_tokens": 95}}, + {"prompt_tokens_details": {"cache_creation_tokens": 95}}, + {"cache_creation_input_tokens": 95}, + ], +) +def test_every_spelling_of_cache_write_tokens_is_read( + logger: DataDogLLMObsLogger, usage_object: dict[str, Any] +) -> None: + """A raw usage dict that bypassed litellm's normalizer can carry any provider's spelling.""" + payload = build(logger, usage_object=usage_object) + + assert payload["metrics"]["cache_write_input_tokens"] == 95.0 + + +def test_a_cache_read_does_not_emit_a_zero_cache_write(logger: DataDogLLMObsLogger) -> None: + """A zero write on every cache-read span would drag Datadog's cache-write average to nothing.""" + payload = build(logger, usage_object={"prompt_tokens_details": {"cached_tokens": 4096}}) + + assert payload["metrics"]["cache_read_input_tokens"] == 4096.0 + assert "cache_write_input_tokens" not in payload["metrics"] + + +def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMObsLogger) -> None: + """An uncached request must not gain zero-valued cache metrics that dilute cache dashboards.""" + payload = build(logger, usage_object={"prompt_tokens_details": None}) + + assert "cache_read_input_tokens" not in payload["metrics"] + assert "cache_write_input_tokens" not in payload["metrics"] + assert "non_cached_input_tokens" not in payload["metrics"] + + +def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]}) + + assert payload["meta"]["tool_definitions"] == [ + { + "name": "get_weather", + "description": "Get current weather for a city", + "schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + + +def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None: + """The Anthropic surface declares tools unwrapped, with input_schema instead of parameters.""" + payload = build( + logger, + model_parameters={"tools": [{"name": "get_weather", "description": "d", "input_schema": {"type": "object"}}]}, + ) + + assert payload["meta"]["tool_definitions"] == [ + {"name": "get_weather", "description": "d", "schema": {"type": "object"}} + ] + + +def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogLLMObsLogger) -> None: + assert "tool_definitions" not in build(logger)["meta"] + + +def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None: + """A truncated argument string is still the only record of what the model tried to call.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"city":'}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == '{"city":' + + +def test_oversized_tool_arguments_ship_unparsed(logger: DataDogLLMObsLogger) -> None: + """ + Decoding attacker-sized compact JSON multiplies memory for a span that is only logging. + + This payload is perfectly valid JSON, so the only reason it arrives as a string is the + size bound; a smaller copy of the same shape comes back as an object below. + """ + oversized = '{"a":"' + "x" * 300_000 + '"}' + + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": oversized}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == oversized + + +def test_valid_arguments_below_the_bound_still_parse(logger: DataDogLLMObsLogger) -> None: + """The size bound must not swallow ordinary arguments; this is the oversized test's control.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"a":"' + "x" * 64 + '"}'}} + ], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == {"a": "x" * 64} + + +def test_a_result_is_named_even_when_its_call_had_unparseable_arguments(logger: DataDogLLMObsLogger) -> None: + """Correlating a result to its call reads ids and names, so bad arguments cannot break linking.""" + payload = build( + logger, + messages=[ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_abc123", "type": "function", "function": {"name": "get_weather", "arguments": "{"}} + ], + }, + {"role": "tool", "tool_call_id": "call_abc123", "content": "18C"}, + ], + ) + + assert payload["meta"]["input"]["messages"][1]["tool_results"] == [ + {"name": "get_weather", "result": "18C", "tool_id": "call_abc123", "type": "function"} + ] + + +def test_deeply_nested_tool_arguments_do_not_drop_the_span(logger: DataDogLLMObsLogger) -> None: + """json.loads raises RecursionError, not JSONDecodeError, on hostile nesting.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "[" * 50_000}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "[" * 50_000 + + +def test_tool_arguments_that_parse_to_a_non_object_stay_a_string(logger: DataDogLLMObsLogger) -> None: + """Datadog types arguments as an object, so a bare JSON scalar must not land there as one.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "42"}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "42" + + +def test_a_tool_without_a_name_is_not_offered_as_a_definition(logger: DataDogLLMObsLogger) -> None: + """A nameless tool cannot be matched to a call, so it is dropped rather than sent blank.""" + payload = build(logger, model_parameters={"tools": [{"function": {"description": "no name"}}, TOOL_DEFINITION]}) + + assert [tool["name"] for tool in payload["meta"]["tool_definitions"]] == ["get_weather"] + + +def test_a_tool_definition_without_a_schema_omits_the_field(logger: DataDogLLMObsLogger) -> None: + """An empty schema object would read as a tool that takes no arguments, which is a different claim.""" + payload = build(logger, model_parameters={"tools": [{"name": "ping", "description": "d"}]}) + + assert payload["meta"]["tool_definitions"] == [{"name": "ping", "description": "d"}] + + +def test_a_non_dict_message_still_reaches_datadog(logger: DataDogLLMObsLogger) -> None: + """Callers can log arbitrary message payloads, and dropping the span over one loses the request.""" + payload = build(logger, messages=["just a bare string"]) + + assert payload["meta"]["input"]["messages"] == [{"input": "just a bare string"}] + + +def test_messages_logged_as_a_bare_string_still_reach_datadog(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages="the whole prompt as one string") + + assert payload["meta"]["input"]["messages"] == [{"input": "the whole prompt as one string"}] + + +def test_non_chat_call_types_log_an_empty_input(logger: DataDogLLMObsLogger) -> None: + """Embedding and image calls carry no messages; fabricating an "None" turn misreads in Datadog.""" + payload = build(logger, messages=None) + + assert payload["meta"]["input"]["messages"] == [] + + +def test_anthropic_tool_blocks_map_to_tool_calls_and_results(logger: DataDogLLMObsLogger) -> None: + """/v1/messages carries tool traffic as content blocks, not OpenAI fields.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": [{"type": "text", "text": "Weather in Tokyo?"}]}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Tokyo"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "18C"}]}, + ], + ) + + assistant, result_turn = payload["meta"]["input"]["messages"][1:3] + assert assistant["tool_calls"] == [ + {"name": "get_weather", "arguments": {"city": "Tokyo"}, "tool_id": "toolu_1", "type": "tool_use"} + ] + assert result_turn["tool_results"] == [ + {"name": "get_weather", "result": "18C", "tool_id": "toolu_1", "type": "function"} + ] + + +def test_content_with_no_text_parts_is_preserved_not_blanked(logger: DataDogLLMObsLogger) -> None: + """A content list the mapper does not understand must ride along, not be erased.""" + blocks = [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}] + payload = build(logger, messages=[{"role": "user", "content": blocks}]) + + assert payload["meta"]["input"]["messages"][0]["content"] == blocks + + +def test_multimodal_content_parts_are_flattened_to_text(logger: DataDogLLMObsLogger) -> None: + """Datadog types Message.content as a string, so content lists collapse to their text.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": [{"type": "text", "text": "describe "}, {"type": "text", "text": "this"}]} + ], + ) + + assert payload["meta"]["input"]["messages"][0]["content"] == "describe this" + + +def test_mapping_input_messages_does_not_mutate_the_shared_payload(logger: DataDogLLMObsLogger) -> None: + """Sibling callbacks read the same messages list, so flattening must not write through it.""" + messages: list[dict[str, Any]] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + kwargs = build_payload(messages=messages) + start = datetime(2026, 9, 1, 12, 0, 0) + + logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=1)) + + assert messages[0]["content"] == [{"type": "text", "text": "hi"}] + + +def test_reasoning_content_survives_the_mapping(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + response_message={"role": "assistant", "content": "answer", "reasoning_content": "thinking"}, + ) + + assert payload["meta"]["output"]["messages"][0]["reasoning_content"] == "thinking" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index ca628aa3405..99d706a9c44 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -541,6 +541,74 @@ def test_llm_call_adapter_extracts_cache_tokens_from_usage_object(): assert data.usage.cache_read_input_tokens == 3 +def test_llm_call_adapter_normalizes_nested_cache_tokens(): + cases: Final = ( + ({"prompt_tokens_details": {"cached_tokens": 3}}, 3, None), + ({"prompt_cache_hit_tokens": 11}, 11, None), + ({"prompt_tokens_details": {"cache_write_tokens": 7}}, None, 7), + ({"prompt_tokens_details": {"cache_creation_tokens": 13}}, None, 13), + ({"prompt_tokens_details": {"cache_creation_input_tokens": 17}}, None, 17), + ) + for usage_object, expected_read, expected_creation in cases: + case_payload = _sample_payload(metadata={"usage_object": usage_object}) + data = LLMCallSpanData.from_standard_logging_payload(case_payload) + assert data.usage.cache_read_input_tokens == expected_read + assert data.usage.cache_creation_input_tokens == expected_creation + + +def test_llm_call_adapter_prefers_nested_count_over_zero_top_level(): + payload = _sample_payload( + metadata={ + "usage_object": { + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens == 5 + assert data.usage.cache_creation_input_tokens == 7 + + +def test_llm_call_adapter_ignores_invalid_cache_values_before_valid_fallbacks(): + payload = _sample_payload( + metadata={ + "usage_object": { + "cache_read_input_tokens": -1, + "cache_creation_input_tokens": "5.0", + "prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens == 5 + assert data.usage.cache_creation_input_tokens == 7 + + +def test_llm_call_adapter_ignores_non_finite_cache_values(): + payload = _sample_payload( + metadata={ + "usage_object": { + "prompt_tokens_details": {"cached_tokens": float("nan")}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens is None + + +def test_llm_call_adapter_preserves_explicit_zero_and_omits_missing_cache_tokens(): + for usage_object, expected_read, expected_creation in ( + ({"prompt_tokens_details": {"cached_tokens": 0}}, 0, None), + ({}, None, None), + ): + case_payload = _sample_payload(metadata={"usage_object": usage_object}) + data = LLMCallSpanData.from_standard_logging_payload(case_payload) + assert data.usage.cache_read_input_tokens == expected_read + assert data.usage.cache_creation_input_tokens == expected_creation + + def test_llm_call_adapter_cache_tokens_none_without_usage_object(): data = LLMCallSpanData.from_standard_logging_payload(_sample_payload()) assert data.usage.cache_creation_input_tokens is None diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d978eb48c12..7d70b9a8862 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2237,3 +2237,202 @@ class TestRecordsOwnGuardrailInformation: ) assert _guardrail_entries(request_data) == [] + + +class _ApplyOnlyObserver(CustomGuardrail): + """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" + + def __init__(self, block: bool = False): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__(guardrail_name="apply-only-observer", event_hook=GuardrailEventHooks.logging_only) + self.block = block + self.calls: list = [] + + @log_guardrail_information + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + from fastapi import HTTPException + + self.calls.append((input_type, list(inputs.get("texts") or []))) + if self.block: + raise HTTPException(status_code=400, detail={"error": "flagged"}) + return GenericGuardrailAPIInputs(texts=["[MASKED]" for _ in inputs.get("texts") or []]) + + +def _logged_call(messages: list | str) -> tuple[dict, object]: + from litellm.types.utils import Choices, Message, ModelResponse + + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]) + kwargs = { + "model": "gpt-5.4-mini", + "messages": messages, + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {"user_api_key_user_id": "u1"}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + return kwargs, response + + +class TestLoggingOnlyApplyGuardrail: + """LIT-4876 regression: a guardrail in mode logging_only that implements only + apply_guardrail must still run against the logged request and response and + record guardrail_information, instead of inheriting the CustomLogger no-op.""" + + @pytest.mark.asyncio + async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): + guardrail = _ApplyOnlyObserver() + messages = [{"role": "user", "content": "hello there"}] + kwargs, response = _logged_call(messages) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + assert out_kwargs["messages"] == [{"role": "user", "content": "hello there"}] + assert out_response.choices[0].message.content == "general kenobi" + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_name"] for e in entries] == ["apply-only-observer", "apply-only-observer"] + assert {e["guardrail_mode"] for e in entries} == {"logging_only"} + assert {e["guardrail_status"] for e in entries} == {"success"} + assert "standard_logging_guardrail_information" not in kwargs["litellm_params"]["metadata"] + assert kwargs["standard_logging_object"] == {"guardrail_information": None} + + @pytest.mark.asyncio + async def test_appends_to_pre_call_verdicts_without_duplicating_them(self): + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + pre_call_entry = {"guardrail_name": "pii-blocker", "guardrail_mode": "pre_call", "guardrail_status": "success"} + kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] = [pre_call_entry] + kwargs["standard_logging_object"]["guardrail_information"] = [pre_call_entry] + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_name"] for e in entries] == ["pii-blocker", "apply-only-observer", "apply-only-observer"] + assert kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] == [pre_call_entry] + + @pytest.mark.asyncio + async def test_request_copy_failure_is_swallowed(self): + import threading + + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there", "lock": threading.Lock()}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [] + assert out_kwargs is kwargs + assert out_response is response + + @pytest.mark.asyncio + async def test_block_verdict_is_recorded_without_raising(self): + guardrail = _ApplyOnlyObserver(block=True) + kwargs, response = _logged_call([{"role": "user", "content": "flagged content"}]) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [("request", ["flagged content"])] + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["guardrail_intervened"] + + @pytest.mark.asyncio + async def test_call_type_without_translation_is_skipped(self): + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.amoderation.value) + + assert guardrail.calls == [] + assert out_kwargs["standard_logging_object"]["guardrail_information"] is None + + @pytest.mark.asyncio + async def test_aembedding_scans_logged_input(self): + from litellm.types.utils import EmbeddingResponse + + guardrail = _ApplyOnlyObserver() + kwargs, _ = _logged_call("hello there") + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.aembedding.value) + + assert guardrail.calls == [("request", ["hello there"])] + assert out_kwargs["messages"] == "hello there" + assert out_response is response + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success"] + + @pytest.mark.asyncio + async def test_native_lifecycle_hook_guardrail_is_left_alone(self): + class _NativeHooks(_ApplyOnlyObserver): + use_native_lifecycle_hooks = True + + guardrail = _NativeHooks() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [] + assert out_kwargs is kwargs + assert out_response is response + + @pytest.mark.asyncio + async def test_aresponses_scans_logged_messages_when_input_is_cleared(self): + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _ApplyOnlyObserver() + kwargs, _ = _logged_call([{"role": "user", "content": "hello there"}]) + kwargs["input"] = None + response = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="gpt-5.4-mini", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "general kenobi"}], + } + ], + ) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.aresponses.value) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success", "success"] + + @pytest.mark.asyncio + async def test_async_success_handler_records_verdict_in_standard_logging_object(self): + import datetime as dt + + from litellm.litellm_core_utils.litellm_logging import Logging + + guardrail = _ApplyOnlyObserver() + guardrail.default_on = True + messages = [{"role": "user", "content": "hello there"}] + _, response = _logged_call(messages) + logging_obj = Logging( + model="gpt-5.4-mini", + messages=messages, + stream=False, + call_type=CallTypes.acompletion.value, + start_time=dt.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=[guardrail], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}}, optional_params={}, model="gpt-5.4-mini", custom_llm_provider="openai" + ) + + await logging_obj.async_success_handler( + result=response, start_time=dt.datetime.now(), end_time=dt.datetime.now() + ) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + entries = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success", "success"] diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index 029b097cb75..ea661d2ea78 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -93,6 +93,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): logger._increment_token_metrics = MagicMock() logger._increment_remaining_budget_metrics = AsyncMock() logger._set_virtual_key_rate_limit_metrics = MagicMock() + logger._set_key_and_team_rate_limit_metrics = MagicMock() logger._set_latency_metrics = MagicMock() logger.set_llm_deployment_success_metrics = MagicMock() logger._increment_cache_metrics = MagicMock() diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index 9c6d2e018ff..bf1d68c7714 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -13,6 +13,7 @@ Covers two follow-up gaps to the unified rate-limit error work: 429s don't silently break when the new class lands. """ +from collections.abc import Mapping from unittest.mock import MagicMock, patch import pytest @@ -471,3 +472,254 @@ def test_should_ignore_non_int_v3_header_values(bad_value): logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( sys.maxsize ) + + +KEY_AND_TEAM_RATE_LIMIT_METRICS = ( + "litellm_api_key_rate_limit_allowed_metric", + "litellm_api_key_rate_limit_used_metric", + "litellm_team_rate_limit_allowed_metric", + "litellm_team_rate_limit_used_metric", +) + + +def _clear_prometheus_registry() -> None: + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _collected_samples(metric_name: str) -> dict[tuple[tuple[str, str], ...], float]: + from prometheus_client import REGISTRY + + return { + tuple(sorted(sample.labels.items())): sample.value + for metric in REGISTRY.collect() + for sample in metric.samples + if sample.name == metric_name + } + + +def _success_kwargs_with_rate_limit_headers(additional_headers: Mapping[str, object] | None) -> dict[str, object]: + return { + "model": "claude-haiku-4-5", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": 20, + "prompt_tokens": 15, + "completion_tokens": 5, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": "claude-haiku-4-5", + "model_id": "model-123", + "model_group": "anthropic-haiku-4-5", + "api_base": "https://api.anthropic.com", + "custom_llm_provider": "anthropic", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "stream": False, + "response": None, + "model_parameters": None, + "metadata": { + "user_api_key_hash": "key-hash", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "team-id", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": { + "litellm_overhead_time_ms": None, + "additional_headers": additional_headers, + }, + }, + } + + +async def _run_success_event( + additional_headers: Mapping[str, object] | None, logger: PrometheusLogger | None = None +) -> None: + import datetime + + now = datetime.datetime.now() + await (logger or PrometheusLogger()).async_log_success_event( + _success_kwargs_with_rate_limit_headers(additional_headers), None, now, now + ) + + +@pytest.mark.asyncio +async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_headers(): + """ + LIT-1672: the v3 limiter mirrors ``x-ratelimit-{api_key,team}-{limit,remaining}-*`` + into the logging payload. The gauges must expose the configured limit as-is + and the window consumption as ``limit - remaining`` for each key / team + dimension, split by ``rate_limit_type``. + """ + _clear_prometheus_registry() + try: + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 7, + "x-ratelimit-api_key-limit-tokens": 20000, + "x-ratelimit-api_key-remaining-tokens": 19947, + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 47, + "x-ratelimit-team-limit-tokens": 40000, + "x-ratelimit-team-remaining-tokens": 39960, + "x-ratelimit-model_per_key-limit-requests": 5, + "x-ratelimit-model_per_key-remaining-requests": 1, + } + ) + + key_requests = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "requests"), + ) + key_tokens = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "tokens"), + ) + team_requests = ( + ("rate_limit_type", "requests"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + team_tokens = ( + ("rate_limit_type", "tokens"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == { + key_requests: 10, + key_tokens: 20000, + } + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == { + key_requests: 3, + key_tokens: 53, + } + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == { + team_requests: 50, + team_tokens: 40000, + } + assert _collected_samples("litellm_team_rate_limit_used_metric") == { + team_requests: 3, + team_tokens: 40, + } + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_should_emit_only_the_dimensions_the_limiter_enforced(): + """ + A key with only ``rpm_limit`` set and no team limits produces only the + key/requests headers, so no tokens series and no team series may appear + (a phantom 0 or sys.maxsize series would misreport an unlimited dimension). + """ + _clear_prometheus_registry() + try: + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 10, + } + ) + + key_requests = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "requests"), + ) + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {key_requests: 10} + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {key_requests: 0} + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {} + assert _collected_samples("litellm_team_rate_limit_used_metric") == {} + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_a_limit(): + """ + Removing a key's ``rpm_limit`` / ``tpm_limit`` (or a team's ``tpm_limit``) + makes the v3 limiter stop emitting that descriptor's headers on later + requests. The old allowed/used samples must disappear instead of keeping + a limit that no longer exists on the scrape. + """ + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 7, + "x-ratelimit-api_key-limit-tokens": 20000, + "x-ratelimit-api_key-remaining-tokens": 19947, + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 47, + "x-ratelimit-team-limit-tokens": 40000, + "x-ratelimit-team-remaining-tokens": 39960, + }, + logger=logger, + ) + await _run_success_event( + { + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 46, + }, + logger=logger, + ) + + team_requests = ( + ("rate_limit_type", "requests"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {} + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {} + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {team_requests: 50} + assert _collected_samples("litellm_team_rate_limit_used_metric") == {team_requests: 4} + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "additional_headers", + [ + None, + {"x-ratelimit-model_per_key-remaining-requests": 42}, + {"x-ratelimit-api_key-limit-requests": 10}, + {"x-ratelimit-api_key-limit-requests": "10", "x-ratelimit-api_key-remaining-requests": "7"}, + {"x-ratelimit-team-limit-tokens": True, "x-ratelimit-team-remaining-tokens": 5}, + ], +) +async def test_should_emit_no_key_or_team_rate_limit_series_without_a_complete_int_pair( + additional_headers, +): + _clear_prometheus_registry() + try: + await _run_success_event(additional_headers) + + for metric_name in KEY_AND_TEAM_RATE_LIMIT_METRICS: + assert _collected_samples(metric_name) == {}, metric_name + finally: + _clear_prometheus_registry() diff --git a/tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py b/tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py new file mode 100644 index 00000000000..519a13751f1 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_requested_model_cardinality.py @@ -0,0 +1,279 @@ +""" +LIT-6611: every unique client-supplied model name that fails routing used to +mint permanent Prometheus series carrying ``requested_model=""`` on the +proxy request metrics and the deployment metrics, with no eviction. The fix +collapses any requested model the router does not recognize (and no wildcard +pattern matches) into the single ``other`` label bucket, while recognized +names, aliases, and wildcard-matched names keep their own label values. +""" + +import sys +import types +from unittest.mock import patch + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.prometheus import ( + UNRECOGNIZED_REQUESTED_MODEL_LABEL, + PrometheusLogger, +) +from litellm.proxy._types import UserAPIKeyAuth + + +class _ClientSideError(Exception): + status_code = 400 + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + yield + + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +@pytest.fixture +def router(): + return litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "fake-key"}, + }, + ], + model_group_alias={"gpt4o-alias": "gpt-4o-mini"}, + ) + + +@pytest.fixture +def team_router(): + return litellm.Router( + model_list=[ + { + "model_name": "team-internal-gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + "model_info": {"team_id": "team-1", "team_public_model_name": "team-alias-gpt"}, + }, + { + "model_name": "team-internal-bedrock", + "litellm_params": {"model": "openai/*", "api_key": "fake-key"}, + "model_info": {"team_id": "team-1", "team_public_model_name": "team-models/*"}, + }, + ] + ) + + +def _requested_model_values(metric) -> set[str]: + index = metric._labelnames.index("requested_model") + return {sample_key[index] for sample_key in metric._metrics} + + +def _series_count(metric) -> int: + return len(metric._metrics) + + +def _total_value(metric) -> float: + return sum(child._value.get() for child in metric._metrics.values()) + + +async def _fire_proxy_failure(logger: PrometheusLogger, model: str) -> None: + await logger.async_post_call_failure_hook( + request_data={"model": model, "metadata": {}, "proxy_server_request": {}}, + original_exception=_ClientSideError(f"model {model} does not exist"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key-1"), + ) + + +@pytest.mark.asyncio +async def test_unknown_models_collapse_to_one_series_on_proxy_request_metrics(router): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + for index in range(25): + await _fire_proxy_failure(logger, f"agent-typo-{index}") + + for metric in ( + logger.litellm_proxy_failed_requests_metric, + logger.litellm_proxy_total_requests_metric, + ): + assert _requested_model_values(metric) == {UNRECOGNIZED_REQUESTED_MODEL_LABEL} + assert _series_count(metric) == 1 + assert _total_value(metric) == 25 + + +@pytest.mark.asyncio +async def test_known_alias_and_wildcard_models_keep_their_own_labels(router): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await _fire_proxy_failure(logger, "gpt-4o-mini") + await _fire_proxy_failure(logger, "gpt4o-alias") + await _fire_proxy_failure(logger, "openai/gpt-4o-audio-preview") + await _fire_proxy_failure(logger, "agent-typo-hallucinated") + + for metric in ( + logger.litellm_proxy_failed_requests_metric, + logger.litellm_proxy_total_requests_metric, + ): + assert _requested_model_values(metric) == { + "gpt-4o-mini", + "gpt4o-alias", + "openai/gpt-4o-audio-preview", + UNRECOGNIZED_REQUESTED_MODEL_LABEL, + } + + +@pytest.mark.asyncio +async def test_team_alias_and_team_wildcard_models_keep_their_own_labels(team_router): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", team_router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await _fire_proxy_failure(logger, "team-alias-gpt") + await _fire_proxy_failure(logger, "team-models/gpt-4o-audio-preview") + await _fire_proxy_failure(logger, "agent-typo-hallucinated") + + for metric in ( + logger.litellm_proxy_failed_requests_metric, + logger.litellm_proxy_total_requests_metric, + ): + assert _requested_model_values(metric) == { + "team-alias-gpt", + "team-models/gpt-4o-audio-preview", + UNRECOGNIZED_REQUESTED_MODEL_LABEL, + } + + +@pytest.mark.asyncio +async def test_unknown_models_collapse_to_other_when_router_is_unavailable(): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", None, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await _fire_proxy_failure(logger, "agent-typo-no-router") + await _fire_proxy_failure(logger, "gpt-4o-mini") + + assert _requested_model_values(logger.litellm_proxy_failed_requests_metric) == { + UNRECOGNIZED_REQUESTED_MODEL_LABEL + } + + +@pytest.mark.asyncio +async def test_sdk_router_originated_metrics_keep_labels_without_proxy_router(): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", None, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + logger.set_llm_deployment_failure_metrics( + request_kwargs={ + "model": "sdk-deployment-group", + "litellm_params": {"metadata": {}}, + "standard_logging_object": {}, + "exception": _ClientSideError("model does not exist"), + } + ) + await logger.log_failure_fallback_event( + original_model_group="sdk-fallback-group", + kwargs={"model": "sdk-fallback-group", "metadata": {}}, + original_exception=_ClientSideError("upstream unavailable"), + ) + + assert _requested_model_values(logger.litellm_deployment_failure_responses) == {"sdk-deployment-group"} + assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == {"sdk-fallback-group"} + + +@pytest.mark.asyncio +async def test_sdk_fallback_labels_survive_non_import_errors_from_proxy_module(monkeypatch): + logger = PrometheusLogger() + broken_proxy_module = types.ModuleType("litellm.proxy.proxy_server") + + def _raise_value_error(_name: str): + raise ValueError("bad proxy env var") + + broken_proxy_module.__getattr__ = _raise_value_error # test-quality-ok: reproduces a proxy_server import raising non-ImportError, no injection seam + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", broken_proxy_module) # test-quality-ok: reproduces a proxy_server import raising non-ImportError, no injection seam + + await logger.log_failure_fallback_event( + original_model_group="sdk-fallback-group", + kwargs={"model": "sdk-fallback-group", "metadata": {}}, + original_exception=_ClientSideError("upstream unavailable"), + ) + + assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == {"sdk-fallback-group"} + + +def test_unknown_models_collapse_to_one_series_on_deployment_metrics(router): + logger = PrometheusLogger() + + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + for index in range(25): + logger.set_llm_deployment_failure_metrics( + request_kwargs={ + "model": f"agent-typo-{index}", + "litellm_params": {"metadata": {}}, + "standard_logging_object": {}, + "exception": _ClientSideError("model does not exist"), + } + ) + logger.set_llm_deployment_failure_metrics( + request_kwargs={ + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": {}, + "exception": _ClientSideError("all deployments cooling down"), + } + ) + + for metric in ( + logger.litellm_deployment_failure_responses, + logger.litellm_deployment_total_requests, + ): + assert _requested_model_values(metric) == { + UNRECOGNIZED_REQUESTED_MODEL_LABEL, + "gpt-4o-mini", + } + assert _series_count(metric) == 2 + assert _total_value(metric) == 26 + + +@pytest.mark.asyncio +async def test_fallback_event_requested_model_is_bounded(router): + logger = PrometheusLogger() + kwargs = {"model": "gpt-4o-mini", "metadata": {}} + + with patch("litellm.proxy.proxy_server.llm_router", router, create=True): # test-quality-ok: production reads proxy_server.llm_router lazily, no injection seam + await logger.log_failure_fallback_event( + original_model_group="agent-typo-hallucinated", + kwargs=kwargs, + original_exception=_ClientSideError("model does not exist"), + ) + await logger.log_success_fallback_event( + original_model_group="agent-typo-hallucinated", + kwargs=kwargs, + original_exception=_ClientSideError("model does not exist"), + ) + await logger.log_failure_fallback_event( + original_model_group="gpt-4o-mini", + kwargs=kwargs, + original_exception=_ClientSideError("upstream unavailable"), + ) + + assert _requested_model_values(logger.litellm_deployment_failed_fallbacks) == { + UNRECOGNIZED_REQUESTED_MODEL_LABEL, + "gpt-4o-mini", + } + assert _requested_model_values(logger.litellm_deployment_successful_fallbacks) == { + UNRECOGNIZED_REQUESTED_MODEL_LABEL + } diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py index 7e997870852..58b15b79e76 100644 --- a/tests/test_litellm/integrations/test_s3.py +++ b/tests/test_litellm/integrations/test_s3.py @@ -2,26 +2,27 @@ from datetime import datetime from unittest.mock import MagicMock, patch import litellm +from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES from litellm.integrations.s3 import S3Logger TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id" -def _standard_logging_payload() -> dict: +def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict: return { - "id": "chatcmpl-test-id", + "id": response_id, "metadata": {"user_api_key_team_alias": None}, } -def _log_event_kwargs() -> dict: +def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict: return { "litellm_params": {"metadata": {}}, - "standard_logging_object": _standard_logging_payload(), + "standard_logging_object": _standard_logging_payload(response_id), } -def _run_log_event(callback_params: dict) -> MagicMock: +def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock: original = litellm.s3_callback_params litellm.s3_callback_params = callback_params try: @@ -30,8 +31,8 @@ def _run_log_event(callback_params: dict) -> MagicMock: mock_boto3_client.return_value = mock_s3_client logger = S3Logger() logger.log_event( - kwargs=_log_event_kwargs(), - response_obj={}, + kwargs=_log_event_kwargs(response_id), + response_obj={"id": response_id}, start_time=datetime(2026, 7, 30, 12, 0, 0), end_time=datetime(2026, 7, 30, 12, 0, 1), print_verbose=lambda *args, **kwargs: None, @@ -154,3 +155,30 @@ def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): put_object_kwargs = mock_s3_client.put_object.call_args.kwargs assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" assert "SSEKMSKeyId" not in put_object_kwargs + + +def test_put_object_key_and_filename_are_bounded_for_an_oversized_response_id(): + """The sync logger bounds both the key and the Content-Disposition filename.""" + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": "logs"}, + response_id="resp_" + "A" * 1100, + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert len(put_object_kwargs["Key"].encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert put_object_kwargs["Key"].startswith("logs/2026-07-30/time-12-00-00-000000_resp_") + filename = put_object_kwargs["ContentDisposition"].removeprefix('inline; filename="').removesuffix('"') + assert len(filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + +def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shrink(): + """A long configured s3_path survives whole when the id can be shortened instead.""" + long_path = "litellm-prod-logs/" + "t" * 921 + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-west-2", "s3_path": long_path}, + response_id="resp_" + "B" * 100, + ) + + key = mock_s3_client.put_object.call_args.kwargs["Key"] + assert key.startswith(long_path + "/2026-07-30/") + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 51671d5101e..a037284d7c1 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1170,6 +1170,294 @@ def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json") +# -------------------------------------------------------------- +# object keys bounded to S3's 1024 UTF-8 byte limit +# -------------------------------------------------------------- +def _oversized_response_id() -> str: + return "resp_" + "A" * 1100 + + +def test_s3_object_key_at_the_byte_limit_is_left_alone(): + """A key that still fits is left byte-identical.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + fixed_len = len("input/2026-08-24/.json") + file_name = "x" * (MAX_S3_OBJECT_KEY_BYTES - fixed_len) + + key = get_s3_object_key(s3_path="input", prefix="", start_time=start_time, s3_file_name=file_name) + + assert key == f"input/2026-08-24/{file_name}.json" + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def test_s3_object_key_is_bounded_for_oversized_response_id(): + """An oversized Responses API id is shortened to a readable head plus a digest.""" + import hashlib + + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + file_name = f"time-06-18-41-948021_{_oversized_response_id()}" + + key = get_s3_object_key(s3_path="input", prefix="DefaultTeamProd/", start_time=start_time, s3_file_name=file_name) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("input/DefaultTeamProd/2026-08-24/time-06-18-41-948021_resp_") + assert key.endswith(f"_{hashlib.sha256(file_name.encode('utf-8')).hexdigest()}.json") + + +@pytest.mark.parametrize( + "s3_path,prefix", + [ + ("input", ""), + ("a" * 900, ""), + ("input", "team-" + "b" * 900 + "/"), + ("c" * 600, "team-" + "d" * 600 + "/key-" + "e" * 600 + "/"), + # many short segments, so the trim lands exactly on the budget edge + ("", "ssss/" * 200), + ], +) +def test_s3_object_key_is_bounded_for_long_paths_and_aliases(s3_path: str, prefix: str): + """Long paths, team aliases and key aliases stay within the cap.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + key = get_s3_object_key( + s3_path=s3_path, + prefix=prefix, + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.endswith(".json") + assert "/2026-08-24/" in key or key.startswith("2026-08-24/") + assert "/" not in key.rsplit("2026-08-24/", 1)[1] + + +def test_s3_object_key_trimmed_prefixes_stay_distinct_per_operator(): + """Prefixes that differ only past the trim point keep separate folders.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + keys = [ + get_s3_object_key( + s3_path="input", + prefix="team-" + "b" * 1000 + suffix + "/", + start_time=start_time, + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + for suffix in ("-one", "-two") + ] + + assert keys[0] != keys[1] + assert all(key.startswith("input/team-" + "b" * 900) for key in keys) + assert all(len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES for key in keys) + + +def test_s3_object_key_bounded_prefix_never_splits_a_multibyte_character(): + """A multibyte prefix is trimmed on a character boundary.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + s3_path = "\u65e5\u672c\u8a9e" * 200 + + key = get_s3_object_key( + s3_path=s3_path, + prefix="\u30c1\u30fc\u30e0" * 200 + "/", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert len(key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert key.startswith(s3_path[:100]) + assert "\ufffd" not in key + + +def test_s3_object_key_stays_unique_for_ids_sharing_a_head(): + """Ids sharing a visible head still get distinct keys.""" + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + keys = { + get_s3_object_key( + s3_path="input", + prefix="", + start_time=start_time, + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}{suffix}", + ) + for suffix in ("first", "second", "third") + } + + assert len(keys) == 3 + + +def test_s3_object_key_bounding_matches_the_documented_layout(): + """The bounded key is `//_.json`.""" + import hashlib + + from litellm.integrations.s3 import get_s3_object_key + + file_name = f"time-06-18-41-948021_{_oversized_response_id()}" + + key = get_s3_object_key( + s3_path="input", + prefix="team/", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=file_name, + ) + + digest = hashlib.sha256(file_name.encode("utf-8")).hexdigest() + assert key == f"input/team/2026-08-24/{file_name[:64]}_{digest}.json" + + +def test_s3_object_key_keeps_the_configured_prefix_when_only_the_id_overflows(): + """A 940 byte configured prefix survives whole when only the id overflows.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + prefix = "team-" + "b" * 934 + "/" + + key = get_s3_object_key( + s3_path="", + prefix=prefix, + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name=f"time-06-18-41-948021_{_oversized_response_id()}", + ) + + assert key.startswith(prefix + "2026-08-24/") + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def test_s3_object_key_spends_the_whole_budget_when_the_prefix_must_be_trimmed(): + """A trimmed prefix keeps every byte the budget allows, not whole segments.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + s3_path = "p" * 400 + "/" + "q" * 600 + + key = get_s3_object_key( + s3_path=s3_path, + prefix="", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name="time-06-18-41-948021_abc", + ) + + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("p" * 400 + "/" + "q" * 500) + + +def test_s3_object_key_keeps_a_single_segment_path_as_far_as_it_fits(): + """A path with no separator is kept as far as it fits, never dropped to the bucket root.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + from litellm.integrations.s3 import get_s3_object_key + + key = get_s3_object_key( + s3_path="a" * 1050, + prefix="", + start_time=datetime(2026, 8, 24, 6, 18, 41, 948021), + s3_file_name="time-06-18-41-948021_chatcmpl-xyz", + ) + + assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + assert key.startswith("a" * 900) + + +def test_create_s3_batch_logging_element_bounds_key_and_keeps_full_response_id(): + """The batch element bounds the key and keeps the full response id in the payload.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + + logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) + response_id = _oversized_response_id() + payload = StandardLoggingPayload( + id=response_id, + metadata={"user_api_key_team_alias": "DefaultTeamProd", "user_api_key_alias": "prod-key"}, + messages=[], + ) + + result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) + + assert result is not None + assert len(result.s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert result.s3_object_key.startswith("DefaultTeamProd/prod-key/2026-08-24/") + assert result.payload["id"] == response_id + + +def test_s3_object_download_filename_is_bounded_for_oversized_response_id(): + """The Content-Disposition filename is bounded too, or the PUT fails with MetadataTooLarge.""" + from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), _oversized_response_id()) + + assert len(file_name.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + assert file_name.startswith("time-2026-08-24T06-18-41-948021_resp_") + assert file_name.endswith(".json") + + +def test_s3_object_download_filenames_stay_distinct_when_shortened(): + """Shortened filenames stay distinct.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + start_time = datetime(2026, 8, 24, 6, 18, 41, 948021) + file_names = { + get_s3_object_download_filename(start_time, _oversized_response_id() + suffix) + for suffix in ("first", "second", "third") + } + + assert len(file_names) == 3 + + +def test_s3_object_download_filename_short_id_is_unchanged(): + """An ordinary response id keeps the filename it had before.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), "resp_abc123") + + assert file_name == "time-2026-08-24T06-18-41-948021_resp_abc123.json" + + +def test_create_s3_batch_logging_element_bounds_the_download_filename(): + """The batch element carries a bounded Content-Disposition filename.""" + from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + logger = S3Logger() + payload = StandardLoggingPayload(id=_oversized_response_id(), metadata={}, messages=[]) + + result = logger.create_s3_batch_logging_element(datetime(2026, 8, 24, 6, 18, 41, 948021), payload) + + assert result is not None + assert len(result.s3_object_download_filename.encode("utf-8")) <= MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES + + +@pytest.mark.asyncio +async def test_audit_log_object_key_is_bounded_for_a_long_configured_path(): + """Audit log keys are bounded by the same builder.""" + from litellm.constants import MAX_S3_OBJECT_KEY_BYTES + + logger = S3Logger() + logger.s3_path = "audit-archive/" + "z" * 1100 + + await logger.async_log_audit_log_event({"id": "1a4f7bd0-6f1e-4d0a-9b3c-9f2e1d5a7c88"}) + + assert len(logger.log_queue) == 1 + assert len(logger.log_queue[0].s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES + assert logger.log_queue[0].s3_object_key.startswith("audit-archive/" + "z" * 900) + + +def test_s3_object_download_filename_drops_characters_that_break_the_header(): + """A quote or separator in the response id cannot escape the quoted header value.""" + from litellm.integrations.s3 import get_s3_object_download_filename + + file_name = get_s3_object_download_filename(datetime(2026, 8, 24, 6, 18, 41, 948021), 'resp_a"b/c') + + assert file_name == "time-2026-08-24T06-18-41-948021_resp_a_b_c.json" + + # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 1af3dd3f613..5628d69de26 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -65,6 +65,7 @@ def _job_record(job: ActiveShadowEvalJob, target_type="key", target_id="key-hash target_type=target_type, target_id=target_id, router_name=job.router_name, + router_names=job.router_names, direction=job.direction, baseline_model=job.baseline_model, shadow_percentage=job.shadow_percentage, @@ -81,6 +82,7 @@ def _router( shadow_text="shadow answer", judge_json='{"preference": "A", "confidence": 0.9, "reasoning": "x"}', classifier_cost=None, + sibling_router_texts=None, ): """One mock router serving the shadow call first, the judge call second, told apart by the internal-origin stamp rather than the model, since a reverse job's shadow arm names @@ -100,6 +102,15 @@ def _router( decision["classifier_cost"] = classifier_cost kwargs["metadata"]["routing_decision"] = decision return {"choices": [{"message": {"content": shadow_text}}], "usage": {"completion_tokens": 5}} + if sibling_router_texts and kwargs["model"] in sibling_router_texts: + kwargs["metadata"]["routing_decision"] = { + "tier_label": "MEDIUM", + "routed_model": f"{kwargs['model']}-pick", + } + return { + "choices": [{"message": {"content": sibling_router_texts[kwargs["model"]]}}], + "usage": {"completion_tokens": 5}, + } return ModelResponse( model=kwargs["model"], choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": shadow_text}}], @@ -1210,29 +1221,36 @@ class TestJobValidation: {"direction": "reverse"}, {"baseline_model": "baseline-model"}, {"direction": "sideways", "baseline_model": "baseline-model"}, + {"direction": "reverse", "baseline_model": "baseline-model", "router_names": ("a", "b")}, ], - ids=["reverse-without-baseline", "forward-with-baseline", "unknown-direction"], + ids=["reverse-without-baseline", "forward-with-baseline", "unknown-direction", "reverse-with-router-set"], ) def test_unsamplable_shapes_are_rejected(self, overrides): with pytest.raises(ValidationError): _job(**overrides) - def test_shadow_target_follows_direction(self): - assert _job().shadow_target == "my-router" - assert _reverse_job().shadow_target == "baseline-model" + def test_arm_target_follows_direction(self): + assert _job().arm_target("my-router") == "my-router" + assert _reverse_job().arm_target("my-router") == "baseline-model" + + def test_rows_from_before_router_names_carry_their_set_in_router_name(self): + assert _job().arm_router_names == ("my-router",) + assert _job(router_names=("my-router", "alt-router")).arm_router_names == ("my-router", "alt-router") @pytest.mark.asyncio class TestDirection: @pytest.mark.parametrize( - "job,routed_by,sampled", + "job,routed_by,attempt_rows", [ - (_job(), None, True), - (_job(), "my-router", False), - (_job(), "other-router", True), - (_reverse_job(), "my-router", True), - (_reverse_job(), None, False), - (_reverse_job(), "other-router", False), + (_job(), None, 1), + (_job(), "my-router", 0), + (_job(), "other-router", 1), + (_reverse_job(), "my-router", 1), + (_reverse_job(), None, 0), + (_reverse_job(), "other-router", 0), + (_job(router_names=("my-router", "alt-router")), "alt-router", 0), + (_job(router_names=("my-router", "alt-router")), "other-router", 2), ], ids=[ "forward-samples-unrouted", @@ -1241,20 +1259,24 @@ class TestDirection: "reverse-samples-its-own-router", "reverse-skips-unrouted", "reverse-skips-another-router", + "forward-skips-any-candidates-own-traffic", + "forward-multi-samples-once-per-arm", ], ) - async def test_direction_decides_which_traffic_is_sampled(self, job, routed_by, sampled): + async def test_direction_decides_which_traffic_is_sampled(self, job, routed_by, attempt_rows): """The two directions partition the key's traffic: whatever one samples, the other - skips, so a key running both never judges the same turn twice for the same reason.""" + skips, so a key running both never judges the same turn twice for the same reason. + A multi-router job extends the forward skip to every candidate: a request one + candidate served must not be judged as the incumbent against another candidate.""" prisma = _prisma() - logger = _logger(router=_router(), prisma=prisma, jobs=(job,)) + logger = _logger(router=_router(sibling_router_texts={"alt-router": "alt answer"}), prisma=prisma, jobs=(job,)) await logger.async_log_success_event( _success_kwargs(request_metadata=_routed_by(routed_by) if routed_by else {}), RESPONSE, None, None ) await _drain(logger) - assert prisma.db.litellm_shadowevalattempt.create.await_count == int(sampled) + assert prisma.db.litellm_shadowevalattempt.create.await_count == attempt_rows async def test_reverse_duplicates_against_the_baseline_model(self): prisma = _prisma() @@ -1316,6 +1338,134 @@ class TestDirection: assert logger._job_starts == {"forward-job": 1, "reverse-job": 1} +@pytest.mark.asyncio +class TestMultiRouterArms: + async def test_every_arm_judges_the_same_request_and_stamps_its_own_row(self): + """One sampled request, one row per candidate router, both judged against the same + real response: the paired comparison that makes multi-router win rates comparable.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router")), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.001, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert [row["router_name"] for row in rows] == ["my-router", "alt-router"] + assert {row["request_id"] for row in rows} == {"req-1"} + assert [row["shadow_model"] for row in rows] == ["cheap-model", "alt-router-pick"] + assert all(row["outcome"] in ("real", "shadow", "tie") for row in rows) + assert all(row["real_cost"] == 0.001 for row in rows) + + async def test_a_single_router_job_stamps_its_router_on_the_row(self): + prisma = _prisma() + logger = _logger(router=_router(), prisma=prisma) + + await logger._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + row = prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"] + assert row["router_name"] == "my-router" + + async def test_one_arms_failure_never_silences_the_sibling(self): + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + healthy = router.acompletion.side_effect + + async def first_arm_explodes(**kwargs): + if kwargs["model"] == "my-router": + raise RuntimeError("provider exploded") + return await healthy(**kwargs) + + router.acompletion.side_effect = first_arm_explodes + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router")), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert [row["router_name"] for row in rows] == ["my-router", "alt-router"] + assert rows[0]["outcome"] == "error" + assert "provider exploded" in rows[0]["error"] + assert rows[1]["outcome"] in ("real", "shadow", "tie") + + async def test_the_turn_valve_counts_every_arm_a_start_will_write(self): + """max_turns is a row ceiling and one sampled request writes one row per arm, so + admission pre-counts the arms: a two-arm job with two turns of budget admits one + request, not two.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger( + router=router, prisma=prisma, jobs=(_job(router_names=("my-router", "alt-router"), max_turns=2),) + ) + + await logger.async_log_success_event(_success_kwargs(request_id="req-1"), RESPONSE, None, None) + await logger.async_log_success_event(_success_kwargs(request_id="req-2"), RESPONSE, None, None) + await _drain(logger) + + rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.await_args_list] + assert {row["request_id"] for row in rows} == {"req-1"} + assert len(rows) == 2 + + async def test_a_withheld_request_runs_no_arm_and_counts_once(self): + """The budget gates run once per sampled request, before any arm: funnel counters + stay per-request, so coverage math is arm-count independent.""" + prisma = _prisma() + router = _router(sibling_router_texts={"alt-router": "alt answer"}) + logger = _logger(router=router, prisma=prisma) + + await logger._run_shadow_eval( + job=_job(router_names=("my-router", "alt-router"), max_budget=1.0, spend=2.0), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + router.acompletion.assert_not_called() + prisma.db.litellm_shadowevalattempt.create.assert_not_called() + assert logger._test_funnel == [("job-1", "withheld")] + + @pytest.mark.asyncio class TestActiveJobsFailClosed: async def test_a_row_the_sampler_cannot_read_is_dropped_not_guessed(self): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3c4121977de..b7f0ca1efe1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -32,6 +32,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, + _is_off_peak, + _is_within_off_peak_window, calculate_cache_writing_cost, generic_cost_per_token, get_token_type_cost_breakdown, @@ -409,6 +411,377 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): assert prompt_base_cost == 9e-6 +def test_is_within_off_peak_window_same_day(): + from datetime import datetime, timezone + + window = "09:00-17:00" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 8, 59, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 17, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_wraps_midnight(): + from datetime import datetime, timezone + + window = "16:30-00:30" + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 15, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 16, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 0, 30, tzinfo=timezone.utc)) is False + assert _is_within_off_peak_window(window, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_equal_start_and_end_covers_whole_day(): + """An equal start and end is the natural way to spell off-peak all day. It used to take the + non-wrap branch, where start <= now < end can never hold, so it matched nothing and billed at + standard rates around the clock without raising or logging anything.""" + from datetime import datetime, timezone + + for window in ("00:00-00:00", "10:00-10:00"): + for hour in range(24): + assert ( + _is_within_off_peak_window(window, datetime(2026, 1, 1, hour, 0, tzinfo=timezone.utc)) is True + ), f"{window} should cover {hour:02d}:00" + + +def test_is_within_off_peak_window_multiple_windows(): + from datetime import datetime, timezone + + # Providers like DeepSeek V4 have more than one daily peak/off-peak window. + windows = ["01:00-05:00", "13:00-16:00"] + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 3, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 14, 30, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window(windows, datetime(2026, 1, 1, 9, 0, tzinfo=timezone.utc)) is False + # a malformed entry in the list is ignored, valid entries still match + assert _is_within_off_peak_window(["bad", "13:00-16:00"], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is True + assert _is_within_off_peak_window([], datetime(2026, 1, 1, 14, 0, tzinfo=timezone.utc)) is False + + +def test_is_within_off_peak_window_normalizes_timezone_aware_input(): + from datetime import datetime, timedelta, timezone + + # A caller may pass a non-UTC aware datetime; the window is UTC and must be + # evaluated in UTC, not against the caller's wall-clock. 09:00 at UTC+8 is + # 01:00 UTC, inside the 01:00-05:00 window. + tz_plus_8 = timezone(timedelta(hours=8)) + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 9, 0, tzinfo=tz_plus_8)) is True + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 12, 0, tzinfo=tz_plus_8)) is True + # 06:00 at UTC+8 is 22:00 UTC the previous day, outside the window + assert _is_within_off_peak_window("01:00-05:00", datetime(2026, 1, 1, 6, 0, tzinfo=tz_plus_8)) is False + + +def test_is_within_off_peak_window_malformed_returns_false(): + from datetime import datetime, timezone + + now = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + assert _is_within_off_peak_window("not-a-window", now) is False + assert _is_within_off_peak_window("16:30", now) is False + assert _is_within_off_peak_window("25:00-26:00", now) is False + + +def test_is_off_peak_weekday_qualified_windows_deepseek_schedule(): + """DeepSeek since 2026-08-23: peak is 01:00-04:00 and 06:00-10:00 UTC on weekdays only, with + weekends off-peak around the clock. The weekday axis is not a filter on one window set; on + two days of seven the off-peak window becomes the whole day, so the schedule needs two + day-qualified rules. The weekend instants inside would-be peak hours are the ones a + time-only implementation bills wrong.""" + from datetime import datetime, timezone + + deepseek = { + "windows": [ + {"hours_utc": ["00:00-01:00", "04:00-06:00", "10:00-00:00"], "weekdays": [1, 2, 3, 4, 5]}, + {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, + ], + } + peak_instants = [ + datetime(2026, 8, 24, 1, 30, tzinfo=timezone.utc), + datetime(2026, 8, 26, 7, 0, tzinfo=timezone.utc), + datetime(2026, 8, 28, 9, 59, tzinfo=timezone.utc), + ] + off_peak_instants = [ + datetime(2026, 8, 23, 1, 30, tzinfo=timezone.utc), + datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc), + datetime(2026, 8, 30, 8, 0, tzinfo=timezone.utc), + datetime(2026, 8, 26, 5, 0, tzinfo=timezone.utc), + datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc), + datetime(2026, 8, 24, 0, 30, tzinfo=timezone.utc), + ] + for when in peak_instants: + assert _is_off_peak(deepseek, when) is False, f"{when.isoformat()} should bill peak" + for when in off_peak_instants: + assert _is_off_peak(deepseek, when) is True, f"{when.isoformat()} should bill off-peak" + + +def test_is_off_peak_weekday_timezone_reads_vendor_calendar(): + """The UTC and Asia/Shanghai calendars only disagree about the date over 16:00-24:00 UTC, so + a window in that stretch is the one place a vendor-local weekday differs from a UTC one: + 2026-08-28T16:30Z is Friday in UTC but already Saturday in Beijing.""" + from datetime import datetime, timezone + + shanghai_saturday = { + "weekday_timezone": "Asia/Shanghai", + "windows": [{"hours_utc": "16:00-17:00", "weekdays": [6]}], + } + assert _is_off_peak(shanghai_saturday, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(shanghai_saturday, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_weekdays_default_utc_calendar_and_accept_names(): + from datetime import datetime, timezone + + named_weekend = {"windows": [{"hours_utc": "00:00-00:00", "weekdays": ["Sat", "sunday"]}]} + assert _is_off_peak(named_weekend, datetime(2026, 8, 29, 12, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(named_weekend, datetime(2026, 8, 28, 12, 0, tzinfo=timezone.utc)) is False + + utc_friday = {"windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(utc_friday, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(utc_friday, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_naive_current_time_read_as_utc(): + from datetime import datetime + + block = {"windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(block, datetime(2026, 8, 28, 16, 30)) is True + assert _is_off_peak(block, datetime(2026, 8, 29, 16, 30)) is False + + +def test_is_off_peak_invalid_weekday_timezone_falls_back_to_utc(): + from datetime import datetime, timezone + + block = {"weekday_timezone": "Not/AZone", "windows": [{"hours_utc": "16:00-17:00", "weekdays": [5]}]} + assert _is_off_peak(block, datetime(2026, 8, 28, 16, 30, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 29, 16, 30, tzinfo=timezone.utc)) is False + + +def test_is_off_peak_ignores_malformed_weekday_rules(): + from datetime import datetime, timezone + + when = datetime(2026, 8, 29, 12, 0, tzinfo=timezone.utc) + assert _is_off_peak({"windows": [{"hours_utc": "00:00-00:00", "weekdays": []}]}, when) is False + assert _is_off_peak({"windows": [{"hours_utc": "00:00-00:00", "weekdays": [0, 8, "noday", True]}]}, when) is False + assert _is_off_peak({"windows": [{"weekdays": [6]}]}, when) is False + assert _is_off_peak({"windows": [{"hours_utc": 1630}]}, when) is False + assert _is_off_peak({"windows": ["00:00-00:00"]}, when) is False + assert _is_off_peak({"windows": "00:00-00:00"}, when) is False + assert _is_off_peak({"hours_utc": 1630}, when) is False + assert _is_off_peak({}, when) is False + + +def test_is_off_peak_flat_hours_and_windows_are_a_union(): + from datetime import datetime, timezone + + block = { + "hours_utc": "04:00-06:00", + "windows": [{"hours_utc": "00:00-00:00", "weekdays": [7]}], + } + assert _is_off_peak(block, datetime(2026, 8, 28, 5, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 30, 20, 0, tzinfo=timezone.utc)) is True + assert _is_off_peak(block, datetime(2026, 8, 28, 20, 0, tzinfo=timezone.utc)) is False + + +def test_get_token_base_cost_weekend_only_off_peak_rate(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": { + "windows": [ + {"hours_utc": ["00:00-01:00", "04:00-06:00", "10:00-00:00"], "weekdays": [1, 2, 3, 4, 5]}, + {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, + ], + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + saturday_peak_hours = _get_token_base_cost( + model_info, usage, current_time=datetime(2026, 8, 29, 2, 0, tzinfo=timezone.utc) + ) + assert saturday_peak_hours[:2] == (5e-7, 1e-6) + + monday_same_hours = _get_token_base_cost( + model_info, usage, current_time=datetime(2026, 8, 24, 2, 0, tzinfo=timezone.utc) + ) + assert monday_same_hours[:2] == (1e-6, 2e-6) + + +def test_get_token_base_cost_applies_off_peak_pricing(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + }, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + assert off_peak[4] == 5e-8 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 1e-6 + assert peak[1] == 2e-6 + assert peak[4] == 1e-7 + + +def test_get_token_base_cost_non_mapping_off_peak_block_bills_standard_rates(): + """A truthy non-mapping off_peak_pricing value (a bare string or a list in + YAML) must bill standard rates rather than raising, matching how every + other malformed piece of the block behaves. + """ + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + when = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + + for malformed_block in ("16:00-19:00", ["16:00-19:00"], 5e-7, True): + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": malformed_block, + }, + ) + result = _get_token_base_cost(model_info, usage, current_time=when) + assert result[0] == 1e-6 + assert result[1] == 2e-6 + + +def test_get_token_base_cost_off_peak_falls_back_to_standard_when_unset(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + + result = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert result[0] == 5e-7 + assert result[1] == 2e-6 + + +def test_get_token_base_cost_off_peak_wins_over_threshold(): + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_token_above_200k_tokens": 3e-6, + "output_cost_per_token_above_200k_tokens": 4e-6, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + }, + ) + usage = Usage(prompt_tokens=250000, completion_tokens=250000, total_tokens=500000) + + off_peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert off_peak[0] == 5e-7 + assert off_peak[1] == 1e-6 + + peak = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert peak[0] == 3e-6 + assert peak[1] == 4e-6 + + +def test_get_model_info_propagates_off_peak_fields(): + model_name = "test-off-peak-model" + off_peak_pricing = { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + "cache_read_input_token_cost": 5e-8, + } + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "off_peak_pricing": off_peak_pricing, + } + } + ) + info = litellm.get_model_info(model=model_name) + assert info["off_peak_pricing"] == off_peak_pricing + + +def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): + """Tiered pricing resolves base rates on its own path and returns early, so off-peak has to + be applied there too or a model carrying both would silently bill the tier rate all day.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-tiered" + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "tiered_pricing": [ + {"range": [0, 128000], "input_cost_per_token": 3e-6, "output_cost_per_token": 6e-6}, + ], + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "input_cost_per_token": 5e-7, + "output_cost_per_token": 1e-6, + }, + } + } + ) + info = litellm.get_model_info(model=model_name) + usage = Usage(prompt_tokens=1_000, completion_tokens=100, total_tokens=1_100) + + inside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)) + assert inside[:2] == (5e-7, 1e-6) + + outside = _get_token_base_cost(info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert outside[:2] == (3e-6, 6e-6) + + def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" @@ -1149,7 +1522,7 @@ def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map): sol = litellm.model_cost["gpt-5.6-sol"] cost_fields = sorted(field for field in sol if "cost" in field) - assert len(cost_fields) == 23 + assert len(cost_fields) == 27 for field in cost_fields: assert alias.get(field) == sol.get(field), field @@ -3666,8 +4039,8 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m ) assert fast == priority - assert fast[0] == pytest.approx(300_000 * 8e-06, rel=1e-9) - assert fast[1] == pytest.approx(1_000 * 3e-05, rel=1e-9) + assert fast[0] == pytest.approx(300_000 * 1.6e-05, rel=1e-9) + assert fast[1] == pytest.approx(1_000 * 6e-05, rel=1e-9) def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map): @@ -3827,6 +4200,86 @@ def test_generic_cost_per_token_gemini_37_flash(_local_model_cost_map): assert completion_cost == pytest.approx(0.001875) +GEMINI_38_FLASH_LAUNCH_PRICING = [ + ("gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), + ("gemini/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), + ("vertex_ai/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), +] + + +@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_38_FLASH_LAUNCH_PRICING) +def test_gemini_38_flash_launch_pricing(model, input_cost, output_cost, cache_read_cost, _local_model_cost_map): + model_cost_map = litellm.model_cost[model] + assert model_cost_map["input_cost_per_token"] == input_cost + assert model_cost_map["output_cost_per_token"] == output_cost + assert model_cost_map["output_cost_per_reasoning_token"] == output_cost + assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost + assert model_cost_map["mode"] == "chat" + assert model_cost_map["supports_reasoning"] is True + assert model_cost_map["supports_function_calling"] is True + assert model_cost_map["max_input_tokens"] == 1048576 + + +GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH = ( + "input_cost_per_token", + "output_cost_per_token", + "output_cost_per_reasoning_token", + "cache_read_input_token_cost", + "input_cost_per_token_batches", + "output_cost_per_token_batches", + "input_cost_per_token_flex", + "output_cost_per_token_flex", + "cache_read_input_token_cost_flex", + "input_cost_per_token_priority", + "output_cost_per_token_priority", + "cache_read_input_token_cost_priority", + "search_context_cost_per_query", + "google_maps_grounding_cost_per_query", + "prompt_cache_min_tokens", + "max_input_tokens", + "max_output_tokens", + "supports_reasoning", + "supports_function_calling", + "supports_prompt_caching", + "supports_vision", + "supports_pdf_input", + "supports_audio_input", + "supports_video_input", + "supports_response_schema", + "supports_tool_choice", + "supports_web_search", + "supports_url_context", +) + + +@pytest.mark.parametrize("prefix", ["", "gemini/", "vertex_ai/"]) +def test_gemini_38_flash_matches_37_flash_promotional_pricing(prefix, _local_model_cost_map): + new_model = litellm.model_cost[f"{prefix}gemini-3.8-flash"] + old_model = litellm.model_cost[f"{prefix}gemini-3.7-flash"] + for field in GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH: + assert new_model[field] == old_model[field], field + + +def test_generic_cost_per_token_gemini_38_flash(_local_model_cost_map): + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=200, + text_tokens=300, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1000), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model="gemini-3.8-flash", + usage=usage, + custom_llm_provider="gemini", + ) + assert prompt_cost == pytest.approx(0.00075) + assert completion_cost == pytest.approx(0.001875) + + def test_grok_46_launch_pricing(_local_model_cost_map): model_cost_map = litellm.model_cost["xai/grok-4.6"] assert model_cost_map["input_cost_per_token"] == 2e-06 diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 9ab66d55f3f..c037f928593 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1435,6 +1435,80 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot +class TestToolWithFlattenedParameters: + def _anyof_tool(self): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def test_flattens_anyof_parameters_into_new_tool(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + tool = self._anyof_tool() + result = tool_with_flattened_parameters(tool) + + assert result is not tool + parameters = result["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert result["function"]["name"] == "automation_update" + assert tool == self._anyof_tool() + + def test_clean_parameters_return_the_same_tool_object(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + + assert tool_with_flattened_parameters(tool) is tool + + @pytest.mark.parametrize( + "tool", + [ + {"type": "function"}, + {"type": "function", "function": "not-a-dict"}, + {"type": "function", "function": {"name": "no_params"}}, + {"type": "function", "function": {"name": "bad_params", "parameters": "not-a-dict"}}, + ], + ) + def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + assert tool_with_flattened_parameters(tool) is tool + + class TestRequestContainsImageContent: """One detector for every dialect that reaches pre-routing hooks untranslated.""" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 72d26f31c60..dd2d45f00c6 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2932,6 +2932,28 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) +def test_add_cache_point_tool_block_stands_down_for_model_without_prompt_caching(monkeypatch): + """A tool carrying cache_control must not become a cachePoint for a Bedrock model + whose cost-map entry lacks prompt caching support, since Bedrock rejects the whole + request. An unmapped id keeps emitting so ARN deployments do not lose caching.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + add_cache_point_tool_block, + ) + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + tool = {"cache_control": {"type": "ephemeral"}} + + assert add_cache_point_tool_block(tool, model="nvidia.nemotron-super-3-120b") is None + assert add_cache_point_tool_block(tool, model="us.nvidia.nemotron-super-3-120b") is None + assert add_cache_point_tool_block( + tool, model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123" + ) == {"cachePoint": {"type": "default"}} + assert add_cache_point_tool_block(tool, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0") == { + "cachePoint": {"type": "default"} + } + + def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(monkeypatch): """ End-to-end: _bedrock_tools_pt should produce cachePoint blocks with ttl @@ -3627,3 +3649,67 @@ def test_convert_gemini_tool_call_result_answers_tool_reference_only_result(): ) assert result == {"function_response": {"name": "ToolSearch", "response": {"content": ""}}} + + +def test_convert_to_anthropic_tool_invoke_degrades_unpaired_server_tool_use(): + """A replayed srvtoolu_ call whose server tool result is not available + (e.g. the Responses bridge replays items without provider_specific_fields) + must become a plain client tool_use so the client's tool_result can pair + with it. A dangling server_tool_use makes Anthropic 400 the request with + "unexpected `tool_use_id` found in `tool_result` blocks".""" + from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_tool_invoke + + result = convert_to_anthropic_tool_invoke( + tool_calls=[ + { + "id": "srvtoolu_01Unpaired", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "zig version"}'}, + } + ], + web_search_results=None, + tool_results=None, + ) + + assert result == [ + { + "type": "tool_use", + "id": "srvtoolu_01Unpaired", + "name": "web_search", + "input": {"query": "zig version"}, + } + ] + + +def test_convert_to_anthropic_tool_invoke_keeps_paired_server_tool_use(): + """When the paired server tool result is available, the srvtoolu_ call is + still reconstructed as server_tool_use followed by its result block.""" + from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_tool_invoke + + server_result = { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01Paired", + "content": [{"type": "web_search_result", "url": "https://ziglang.org", "title": "Zig"}], + } + + result = convert_to_anthropic_tool_invoke( + tool_calls=[ + { + "id": "srvtoolu_01Paired", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "zig version"}'}, + } + ], + web_search_results=[server_result], + tool_results=None, + ) + + assert result == [ + { + "type": "server_tool_use", + "id": "srvtoolu_01Paired", + "name": "web_search", + "input": {"query": "zig version"}, + }, + server_result, + ] diff --git a/tests/test_litellm/litellm_core_utils/test_audio_utils.py b/tests/test_litellm/litellm_core_utils/test_audio_utils.py index 0e8176fffce..155f6680416 100644 --- a/tests/test_litellm/litellm_core_utils/test_audio_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_audio_utils.py @@ -347,3 +347,65 @@ class TestNormalizeTranscriptionLanguageToBcp47: ) assert normalize_transcription_language_to_bcp47(language) == expected + + +class TestResolveSpeechMediaType: + @pytest.mark.parametrize( + ("upstream_content_type", "response_format", "expected"), + [ + ("audio/wav", None, "audio/wav"), + ("AUDIO/WAV", None, "audio/wav"), + ("audio/flac; charset=binary", "mp3", "audio/flac"), + ("application/json", "flac", "audio/flac"), + ("application/octet-stream", "pcm", "audio/pcm"), + (None, "wav", "audio/wav"), + (None, "WAV", "audio/wav"), + (None, "opus", "audio/opus"), + (None, "aac", "audio/aac"), + (None, "mp3", "audio/mpeg"), + (None, "mp4", "audio/mpeg"), + (None, "bogus", "audio/mpeg"), + (None, None, "audio/mpeg"), + ("", None, "audio/mpeg"), + ], + ) + def test_resolution(self, upstream_content_type, response_format, expected): + from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type + + resolved = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=response_format, + ) + assert resolved == expected + + +class TestSpeechMediaTypeFromAudioBytes: + @pytest.mark.parametrize( + ("audio", "expected"), + [ + (b"RIFF\x24\x00\x00\x00WAVEfmt ", "audio/wav"), + (b"fLaC\x00\x00\x00\x22", "audio/flac"), + (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"), + (b"OggS" + b"\x00" * 24 + b"\x01vorbis", "audio/ogg"), + (b"ID3\x04\x00\x00\x00\x00\x00\x00", "audio/mpeg"), + (b"\xff\xfb\x90\x64", "audio/mpeg"), + (b"\xff\xf3\x80\x00", "audio/mpeg"), + (b"\xff\xf1\x50\x80", "audio/aac"), + (b"\xff\xf9\x50\x80", "audio/aac"), + (b"RIFF\x24\x00\x00\x00AVI LIST", None), + (b"\xff\xff\xff\xff\xff\xff", None), + (b"\xff\xfb\xf0\x00", None), + (b"\xff\xfb\x9c\x00", None), + (b"\xff\xeb\x90\x00", None), + (b"\xff\xf1\xf4\x80", None), + (b"\xff\x00\x00\x00", None), + (b"\x00\x01\x02\x03\x04\x05", None), + (b"\xff\xfb", None), + (b"\xff", None), + (b"", None), + ], + ) + def test_sniffing(self, audio, expected): + from litellm.litellm_core_utils.audio_utils.utils import speech_media_type_from_audio_bytes + + assert speech_media_type_from_audio_bytes(audio) == expected diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 0587628e2fe..b1e8163b91d 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -366,7 +366,7 @@ def test_shipped_rules_stack_adaptive_and_mid_conversation_flags(shipped_cost_ma def test_shipped_rules_flag_unmapped_fable_as_always_on_thinking(shipped_cost_map): """An unmapped Fable/Mythos id picks up ``thinking_always_on`` from the claude-always-on-thinking rule, while other unmapped Claudes stay unflagged.""" - model = "claude-fable-5-1" + model = "claude-fable-6-1" assert model not in litellm.model_cost info = litellm.get_model_info(model, custom_llm_provider="anthropic") assert info["thinking_always_on"] is True @@ -419,7 +419,7 @@ def test_shipped_rules_cover_new_families_like_fable_at_5_plus(shipped_cost_map) """Both version gates accept any claude-- id at major 5 or higher, bare major or major-minor, so a new family shaped like claude-fable-5 gets adaptive thinking and mid-conversation system support without a cost-map entry.""" - model = "claude-fable-5-1" + model = "claude-fable-6-1" assert model not in litellm.model_cost info = litellm.get_model_info(model, custom_llm_provider="anthropic") assert info["supports_mid_conversation_system"] is True diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 8c0e8ee5d02..a374e03d1c7 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -256,14 +256,12 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch): from litellm.litellm_core_utils import get_model_cost_map as module monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) - monkeypatch.setattr( - module.GetModelCostMap, - "fetch_remote_model_cost_map", - staticmethod(lambda url, timeout=5: _load_root_cost_map()), + client, _calls = _mock_client( + [httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client ) before = datetime.now(timezone.utc) - module.get_model_cost_map(url="https://example.invalid/cost_map.json") + module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) loaded_at = module.get_model_cost_map_loaded_at() assert loaded_at is not None @@ -308,7 +306,7 @@ def _unset_local_cost_map_env(monkeypatch): monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) -def _mock_client(outcomes): +def _mock_client(outcomes, client_cls=httpx.AsyncClient): """httpx client over a MockTransport serving one outcome per request; an exception instance is raised.""" calls = {"count": 0} @@ -320,7 +318,7 @@ def _mock_client(outcomes): raise outcome return outcome - return httpx.AsyncClient(transport=httpx.MockTransport(handler)), calls + return client_cls(transport=httpx.MockTransport(handler)), calls @pytest.mark.asyncio @@ -450,3 +448,97 @@ async def test_refetch_respects_local_env_override(monkeypatch): ) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 + + +# --------------------------------------------------------------------------- +# get_model_cost_map: the boot-time load retries transient failures like a reload does +# --------------------------------------------------------------------------- + +from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + get_model_cost_map_source_info, +) + + +class _SyncSleepRecorder: + """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" + + def __init__(self): + self.waits = [] + + def __call__(self, seconds: float) -> None: + self.waits.append(seconds) + + +def test_boot_load_retries_transient_failures_instead_of_falling_back(): + """A refused connection then a 503 at pod boot used to pin the process to the bundled + backup for its lifetime; both are transient and must be retried before giving up.""" + client, calls = _mock_client( + [ + httpx.ConnectError("connection refused"), + httpx.Response(503), + httpx.Response(200, content=_real_map_bytes()), + ], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + + assert calls["count"] == 3 + assert len(sleeper.waits) == 2 + assert 2.0 <= sleeper.waits[0] < 3.0 + assert 4.0 <= sleeper.waits[1] < 5.0 + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + + +def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts(): + """An outage longer than the retry budget still ends on the bundled backup, and the + recorded fallback reason says how many attempts were spent so operators can tell.""" + client, calls = _mock_client( + [httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client + ) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + + assert calls["count"] == 3 + assert sleeper.waits == [7.0, 7.0] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert "after 3 attempts" in source["fallback_reason"] + assert len(cost_map) > 100 + + +def test_boot_load_does_not_retry_permanent_failures(): + """A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup.""" + client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert get_model_cost_map_source_info()["source"] == "local" + + get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0)) + assert sleeper.waits == [] + assert get_model_cost_map_source_info()["source"] == "local" + + +def test_boot_load_respects_local_env_override(monkeypatch): + """LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + def _fail(request): + raise AssertionError("no HTTP request should be made when local map is forced") + + cost_map = get_model_cost_map( + url=_URL, + sleep=_SyncSleepRecorder(), + client=httpx.Client(transport=httpx.MockTransport(_fail)), + ) + assert len(cost_map) > 100 + assert get_model_cost_map_source_info()["is_env_forced"] is True diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py index cb4e72ab3ad..722818598af 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -188,6 +188,8 @@ class TestDeclaredAuthenticatingProvider: ("gpt-4o", "github_copilot", "github_copilot"), ("openai/gpt-4o", None, None), ("gpt-4o", "openai", None), + ("github_copilot", None, None), + ("chatgpt", None, None), ], ) def test_names_only_the_providers_whose_resolution_authenticates(self, model, provider, expected): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c7328adb0b3..366f61ded49 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5479,6 +5479,97 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj): assert "key=*****" in raw_api_base +def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]): + import datetime + + obj = LitellmLogging( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="slot-leak-test", + function_id="slot-leak-test", + ) + obj.model_call_details["litellm_params"] = {"metadata": {}} + return patch.object(obj, "get_combined_callback_list", return_value=callbacks), obj + + +def _assembled_stream_result(): + response = ModelResponse() + response.choices[0].message.content = "hello" + return response + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_logging_hook_failure(): + """Regression for leaked max_parallel_requests slots: a raising + async_logging_hook must not abort the success-callback loop that + releases the rate-limiter slot.""" + broken = CustomLogger() + broken.async_logging_hook = AsyncMock(side_effect=RuntimeError("broken stream payload")) + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([broken, releasing]) + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_cost_calculation_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + assert logging_obj.model_call_details["response_cost"] is None + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_standard_logging_payload_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + assert logging_obj.model_call_details.get("standard_logging_object") is None + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failure(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + skipping = CustomGuardrail(guardrail_name="skipping-guardrail") + skipping.should_run_guardrail = MagicMock(return_value=False) + skipping.async_logging_hook = AsyncMock() + raising = CustomGuardrail(guardrail_name="raising-guardrail") + raising.should_run_guardrail = MagicMock(return_value=True) + raising.async_logging_hook = AsyncMock(side_effect=RuntimeError("guardrail hook failed")) + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([skipping, raising, releasing]) + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + skipping.async_logging_hook.assert_not_awaited() + raising.async_logging_hook.assert_awaited_once() + releasing.async_log_success_event.assert_awaited_once() + + def _resolve(custom_llm_provider, litellm_params, optional_params, model): from litellm.litellm_core_utils.litellm_logging import ( _resolve_vertex_location_for_cost, diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 8ac050a04f9..bacbcbf132b 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -592,6 +592,59 @@ def test_stream_chunk_builder_litellm_usage_chunks(): assert usage.total_tokens == 77 +def test_calculate_usage_honors_openai_sdk_completion_usage_chunks(): + from openai.types.completion_usage import CompletionUsage + + content_chunk = ModelResponseStream( + id="chatcmpl-sdk-usage-1", + created=1745513206, + model="mantle-claude", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + provider_specific_fields=None, + content="ok", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + ) + usage_chunk = ModelResponseStream( + id="chatcmpl-sdk-usage-1", + created=1745513207, + model="mantle-claude", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[], + provider_specific_fields=None, + stream_options={"include_usage": True}, + ) + usage_chunk.usage = CompletionUsage( + prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704 + ) + assert type(usage_chunk.usage) is CompletionUsage + + chunks = [content_chunk, usage_chunk] + usage = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, model="mantle-claude", completion_output="" + ) + + assert usage.prompt_tokens == 20 + assert usage.completion_tokens == 60 + assert usage.total_tokens == 80 + assert getattr(usage, "cost", None) == pytest.approx(0.000704) + + def test_get_model_from_chunks_azure_model_router(): """ Test that _get_model_from_chunks finds the actual model from Azure Model Router chunks. diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7f54fbfb4c2..4c99bce2b0f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4692,3 +4692,82 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}]) assert assembled is not None assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" + + +class TestStableStreamingResponseId: + """ + All chunks of one streamed response must share the same top-level id + (OpenAI streaming contract). Providers streaming via GenericStreamingChunk + (e.g. GigaChat) do not propagate an upstream response id, so + CustomStreamWrapper must pin the id from the first chunk it creates, + mirroring the existing `created` pinning (issue #11437). + + Clients such as goose merge streamed deltas into one assistant message by + chunk id; per-chunk ids split a single reply into many messages. + """ + + def test_generic_chunks_share_one_id(self): + def _generic_chunks(): + return iter( + [ + { + "text": "Hello", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": " world", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, + }, + { + "text": "", + "tool_use": None, + "is_finished": True, + "finish_reason": "stop", + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + }, + "index": 0, + }, + ] + ) + + wrapper = CustomStreamWrapper( + completion_stream=_generic_chunks(), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + ids = [chunk.id for chunk in wrapper if chunk.id] + assert ids, "no chunks emitted" + assert len(set(ids)) == 1, f"chunk ids differ across one stream: {ids}" + + def test_creator_pins_id_from_first_chunk(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + first = wrapper.model_response_creator() + assert wrapper.response_id == first.id + assert wrapper.model_response_creator().id == first.id + + def test_provider_supplied_id_still_wins(self): + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gigachat/GigaChat-2-Max", + logging_obj=MagicMock(), + custom_llm_provider="gigachat", + ) + wrapper.response_id = "chatcmpl-from-provider" + assert wrapper.model_response_creator().id == "chatcmpl-from-provider" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index bd750a47f63..043537f8c1f 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -3,11 +3,13 @@ import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, @@ -46,6 +48,43 @@ async def test_make_call_passes_logging_obj_to_client_post(): assert call_kwargs.get("logging_obj") is logging_obj +def test_anthropic_completion_does_not_send_deployment_default_limits(): + captured_requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "msg_default_limits", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + try: + litellm.completion( + model="anthropic/claude-3-5-haiku-20241022", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-key", + client=client, + default_api_key_rpm_limit=60, + default_api_key_tpm_limit=5000000, + ) + finally: + client.close() + + request_body = json.loads(captured_requests[0].content) + assert "default_api_key_rpm_limit" not in request_body + assert "default_api_key_tpm_limit" not in request_body + + def test_redacted_thinking_content_block_delta(): chunk = { "type": "content_block_start", diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 25e2c3cda80..0f9f8259bef 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6176,9 +6176,10 @@ def test_is_anthropic_usage_object_rejects_responses_api_usage(): [ # always-on-thinking models reject thinking.type=disabled with a 400 ("claude-fable-5", True), + ("claude-fable-5-1", True), ("claude-mythos-5", True), # unmapped future family member -> claude-always-on-thinking fallback rule - ("claude-fable-5-1", True), + ("claude-fable-6-1", True), # adaptive-capable models that ACCEPT disabled must keep it verbatim ("claude-opus-5", False), ("claude-sonnet-5", False), @@ -6207,3 +6208,179 @@ def test_disabled_thinking_omitted_only_for_always_on_models( assert "thinking" not in request else: assert request["thinking"] == {"type": "disabled"} + + +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "required"}, {"type": "function", "function": {"name": "get_weather"}}], +) +def test_forced_tool_choice_raises_clean_error_on_fable_5_1_without_drop_params( + local_model_cost_map, tool_choice, monkeypatch +): + """Fable 5.1 400s on tool_choice type any/tool (thinking is always on and a + forced call would skip it); without drop_params the caller gets a clean + client-side 400 that explains the workaround, not a provider error.""" + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="forced tool use"): + config.map_openai_params( + non_default_params={"tool_choice": tool_choice}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + + +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "required"}, {"type": "function", "function": {"name": "get_weather"}}], +) +def test_forced_tool_choice_downgraded_to_auto_on_fable_5_1_with_drop_params( + local_model_cost_map, tool_choice +): + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"tool_choice": tool_choice}, + optional_params={}, + model="claude-fable-5-1", + drop_params=True, + ) + + assert result["tool_choice"] == {"type": "auto"} + + +def test_forced_tool_choice_downgrade_keeps_parallel_tool_calls_flag(local_model_cost_map): + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"tool_choice": "required", "parallel_tool_calls": False}, + optional_params={}, + model="claude-fable-5-1", + drop_params=True, + ) + + assert result["tool_choice"] == {"type": "auto", "disable_parallel_tool_use": True} + + +@pytest.mark.parametrize("tool_choice, expected_type", [("auto", "auto"), ("none", "none")]) +def test_unforced_tool_choice_forwarded_on_fable_5_1( + local_model_cost_map, tool_choice, expected_type, monkeypatch +): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"tool_choice": tool_choice}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + + assert result["tool_choice"]["type"] == expected_type + + +@pytest.mark.parametrize("model", ["claude-fable-5", "claude-opus-5", "claude-sonnet-5"]) +def test_forced_tool_choice_forwarded_on_models_that_support_it( + local_model_cost_map, model, monkeypatch +): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"tool_choice": "required"}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert result["tool_choice"] == {"type": "any"} + + +def test_forced_tool_choice_gating_driven_by_model_map_flag(local_model_cost_map, monkeypatch): + """The gate must read ``supports_forced_tool_use`` from the model map, not + the model name: a flagged entry gates a model whose name says nothing.""" + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False}) + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"tool_choice": "required"}, + optional_params={}, + model="claude-zeta-9", + drop_params=True, + ) + + assert result["tool_choice"] == {"type": "auto"} + + +def test_anthropic_drop_params_keeps_format_only_output_config(monkeypatch): + """``drop_params=True`` must not consume ``output_config.format``: the drop + gate is an effort gate and ``format`` is a structured-output field.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_anthropic_drop_params_reduces_mixed_output_config_to_format(monkeypatch): + """``drop_params=True`` drops the effort key on unsupported models but keeps + ``format`` so structured outputs still reach the provider.""" + monkeypatch.setattr(litellm, "drop_params", True) + config = AnthropicConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"z": {"type": "integer"}}}, + } + + result = config.transform_request( + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_config": {"effort": "low", "format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(local_model_cost_map, monkeypatch): + """Backstop: on the tool-based structured-output path, a model flagged + ``supports_forced_tool_use: false`` must not get the forced response-format + tool_choice the provider would 400 on.""" + monkeypatch.setitem( + litellm.model_cost, + "claude-test-no-forced-tools", + {"litellm_provider": "anthropic", "mode": "chat", "supports_forced_tool_use": False}, + ) + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + } + }, + optional_params={}, + model="claude-test-no-forced-tools", + drop_params=False, + ) + + assert "tools" in result + assert "tool_choice" not in result diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 2cf7cd142d6..4e6b9ed0188 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -11,6 +11,7 @@ sys.path.insert( import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY +from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig from litellm.utils import get_optional_params @@ -195,3 +196,91 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None: assert "presence_penalty" not in mapped assert "logit_bias" not in mapped assert "reasoning_effort" in supported + + +class TestAzureToolSchemaCombinatorFlattening: + """ + Regression tests for LIT-6510: Azure's chat completions validator rejects + tool parameters carrying a top-level anyOf/oneOf/allOf for every model + family, so AzureOpenAIConfig.transform_request must flatten them. + """ + + @staticmethod + def _anyof_tool(): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def _transform(self, config, model, tools): + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + def test_transform_request_flattens_top_level_anyof(self): + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [self._anyof_tool()]) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert request["tools"][0]["function"]["name"] == "automation_update" + + def test_gpt5_config_flattens_via_shared_transform(self): + request = self._transform(AzureOpenAIGPT5Config(), "gpt-5.4-mini", [self._anyof_tool()]) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + + def test_caller_tool_dict_is_not_mutated(self): + tool = self._anyof_tool() + self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + assert tool == self._anyof_tool() + + def test_clean_object_schema_passes_through_as_same_object(self): + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + assert request["tools"][0] is tool + + def test_non_dict_tool_entries_pass_through_unchanged(self): + request = self._transform(AzureOpenAIConfig(), "gpt-4o", ["not-a-tool"]) + assert request["tools"] == ["not-a-tool"] + + def test_request_without_tools_is_unchanged(self): + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"temperature": 0.2}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + assert "tools" not in request + assert request["temperature"] == 0.2 diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py index fc7e94a77ba..202f81f1252 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -23,3 +23,48 @@ async def test_azure_chat_o_series_transformation(): ) print(response) assert response["model"] == "web-interface-o1-mini" + + +def test_azure_o_series_transform_request_flattens_top_level_anyof(): + """Regression test for LIT-6510: the o-series super() chain ends in + OpenAIGPTConfig, whose flatten gate skips provider 'azure', so + AzureOpenAIO1Config must flatten tool schema combinators itself.""" + tool = { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + optional_params = {"tools": [tool]} + + request = AzureOpenAIO1Config().transform_request( + model="o3-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert "anyOf" in tool["function"]["parameters"] + assert optional_params["tools"][0] is tool diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index cea299280f8..41d82e4f960 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -428,30 +428,58 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): def test_output_config_format_converted_for_bedrock_chat_invoke_request(): - """Bedrock Invoke chat path consumes ``output_config.format`` before forwarding.""" + """Bedrock Invoke chat path inlines ``output_config.format`` for models + without native structured-output support and keeps the effort key.""" config = AmazonAnthropicClaudeConfig() schema = { "type": "object", "properties": {"answer": {"type": "string"}}, } - result = config.transform_request( + with patch( # test-quality-ok: pin non-native path + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", + ): + result = config.transform_request( + model="anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "test"}], + optional_params={ + "max_tokens": 100, + "output_config": { + "effort": "xhigh", + "format": {"type": "json_schema", "schema": schema}, + }, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"effort": "xhigh"} + last_content = result["messages"][0]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(): + """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort + for models with native structured-output support (Claude Opus 4.7).""" + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( model="anthropic.claude-opus-4-7", messages=[{"role": "user", "content": "test"}], optional_params={ "max_tokens": 100, - "output_config": { - "effort": "xhigh", - "format": {"type": "json_schema", "schema": schema}, - }, + "output_config": {"effort": "xhigh", "format": schema_format}, }, litellm_params={}, headers={}, ) - assert result.get("output_config") == {"effort": "xhigh"} - last_content = result["messages"][0]["content"] - assert json.loads(last_content[-1]["text"]) == schema + assert result.get("output_config") == {"effort": "xhigh", "format": schema_format} + assert "answer" not in json.dumps(result["messages"]) @pytest.mark.parametrize( @@ -488,7 +516,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( optional_params = {"max_tokens": 100, "output_config": {"effort": "high"}} with patch( - "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = config.transform_request( @@ -499,11 +527,7 @@ def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider( headers={}, ) - mock_supports_factory.assert_called_once_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_once_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -542,3 +566,108 @@ def test_output_format_removed_from_bedrock_invoke_request(): assert ( "output_format" not in result ), f"output_format should be removed for Bedrock Invoke, got keys: {result.keys()}" + + +def test_bedrock_chat_invoke_forwards_output_config_format_natively(local_model_cost_map): + """Regression: ``output_config.format`` is forwarded verbatim on models Bedrock + enforces structured outputs for, instead of being inlined as prompt text.""" + import json + + config = AmazonAnthropicClaudeConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + "required": ["zebra_count"], + "additionalProperties": False, + }, + } + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_chat_invoke_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not eat ``output_config.format`` before the + native-forwarding router runs (Sonnet 4.5 has no effort flags).""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = AmazonAnthropicClaudeConfig().transform_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={"max_tokens": 100, "output_config": {"format": schema_format}}, + litellm_params={}, + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_model_cost_map, monkeypatch): + """``drop_params=True`` on a model without native structured-output support + still reaches the inline-schema fallback instead of losing the schema.""" + import litellm + + monkeypatch.setattr(litellm, "drop_params", True) + schema = {"type": "object", "properties": {"zebra_count": {"type": "integer"}}} + + result = AmazonAnthropicClaudeConfig().transform_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "say hello"}], + optional_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params={}, + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +@pytest.mark.parametrize( + "model", + ["us.anthropic.claude-fable-5-1", "anthropic.claude-fable-5-1"], +) +def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(local_model_cost_map, model): + """Regression: Bedrock rejects both native ``output_config.format`` and forced + tool_choice for Fable 5.1, so invoke must use the tool-based path without a + forced ``tool_choice``.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + } + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "output_format" not in result + assert "tools" in result + assert "tool_choice" not in result diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 8e67a7e3438..21e3239f623 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -96,10 +96,8 @@ def _completion_kwargs(**overrides): return kwargs -def _run(**overrides): - with patch.object( - BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS - ): +def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides): + with patch.object(BedrockConverseLLM, "get_credentials", return_value=credentials): return BedrockConverseLLM().completion(**_completion_kwargs(**overrides)) @@ -360,7 +358,7 @@ async def test_async_completion_logs_pre_call_by_default(): def _sync_client_returning_converse_response(): client = MagicMock() - client.post = lambda **_kwargs: httpx.Response( + client.post.side_effect = lambda **_kwargs: httpx.Response( 200, json=CONVERSE_RESPONSE, request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), @@ -487,3 +485,31 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(): assert response.choices[0].message.content == "hi" assert len(calls["post_call"]) == 1 assert "hi" in calls["post_call"][0]["original_response"] + + +def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monkeypatch): + """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no + credentials at all. Preparing the Rust handoff must not dereference that + None: the bearer token signs the request on its own.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") + client = _sync_client_returning_converse_response() + + response = _run(credentials=None, litellm_params={}, client=client) + + assert response.choices[0].message.content == "hi" + sent_headers = client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"] == "Bearer bedrock-bearer-token" + + +def test_the_rust_opt_in_needs_no_sigv4_principal(): + """The core resolves the bearer token itself, so a bearer-only deployment + keeps its opt-in and the gate sees no aws_* credential keys to sign with.""" + seen = _inject() + + response = _run(credentials=None, api_key="bedrock-bearer-token") + + assert response.choices[0].message.content == "hello from rust" + params = seen["call"][0]["optional_params"] + assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys() + assert params["aws_region_name"] == "us-east-1" + assert seen["call"][0]["api_key"] == "bedrock-bearer-token" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 63f895e1819..70f3153ed7e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5248,6 +5248,84 @@ def test_cache_control_injection_tool_config_drops_ttl_for_unsupported_model(): assert tools[-1] == {"cachePoint": {"type": "default"}} +@pytest.mark.parametrize( + ("model", "expects_cache_points"), + [ + pytest.param("nvidia.nemotron-super-3-120b", False, id="mapped-model-without-prompt-caching"), + pytest.param("us.nvidia.nemotron-super-3-120b", False, id="regional-prefix-resolves-through-base-model"), + pytest.param( + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", False, id="claude-named-but-not-caching-on-bedrock" + ), + pytest.param("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True, id="mapped-model-with-prompt-caching"), + pytest.param( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + True, + id="unmapped-arn-keeps-emitting", + ), + ], +) +def test_cache_points_emitted_only_for_models_that_support_prompt_caching(model, expects_cache_points, monkeypatch): + """Bedrock rejects cachePoint blocks for models without prompt caching support + ("You invoked an unsupported model or your request did not allow prompt caching"), + and clients like Claude Code attach cache_control to every request, so a map-known + model without the capability must not receive them. Unmapped ids (application + inference profile ARNs, models newer than the map) keep emitting so existing + caching setups never silently degrade.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + body = AmazonConverseConfig().transform_request( + model=model, + messages=[ + {"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert ("cachePoint" in json.dumps(body)) is expects_cache_points + assert body["system"][0]["text"] == "sys" + assert body["messages"][0]["content"][0]["text"] == "hi" + + +def test_tool_config_cachepoint_not_placed_or_credited_for_model_without_prompt_caching(monkeypatch): + """The tool_config injection point must stand down with the rest of the cachePoint + emission when the model cannot cache, and spend attribution must not credit the + gateway for a breakpoint that was never placed.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + bucket: dict = {"user_api_key": "sk-test"} + data = AmazonConverseConfig()._transform_request_helper( + model="nvidia.nemotron-super-3-120b", + system_content_blocks=[], + optional_params={ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + "cache_control_injection_points": [{"location": "tool_config"}], + }, + messages=[{"role": "user", "content": "hi"}], + litellm_params={"metadata": bucket, "litellm_metadata": None, "model_info": {"id": "dep-bedrock"}}, + ) + + assert "cachePoint" not in json.dumps(data.get("toolConfig", {})) + assert "litellm_gateway_injected_cache" not in bucket + + def test_translate_response_format_json_schema_still_injects_tool(): """ response_format with an explicit json_schema should still use the @@ -6211,7 +6289,7 @@ def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target) result = _bedrock_converse_messages_pt( messages=_agentic_messages_with_ttl(ttl_target), - model="anthropic.claude-3-5-sonnet-20240620-v1:0", + model="anthropic.claude-3-5-sonnet-20241022-v2:0", llm_provider="bedrock_converse", ) @@ -6449,3 +6527,95 @@ def test_disabled_thinking_omitted_for_always_on_models_converse( assert "thinking" not in additional else: assert additional.get("thinking") == {"type": "disabled"} + +@pytest.mark.parametrize( + "model", + ["anthropic.claude-fable-5-1", "us.anthropic.claude-fable-5-1"], +) +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "function", "function": {"name": "get_weather"}}], +) +def test_forced_tool_choice_downgraded_to_auto_on_fable_5_1_converse( + local_model_cost_map, model, tool_choice +): + config = AmazonConverseConfig() + + result = config.map_tool_choice_values( + model=model, tool_choice=tool_choice, drop_params=True + ) + + assert result == {"auto": {}} + + +@pytest.mark.parametrize( + "tool_choice", + ["required", {"type": "function", "function": {"name": "get_weather"}}], +) +def test_forced_tool_choice_raises_clean_error_on_fable_5_1_converse( + local_model_cost_map, tool_choice, monkeypatch +): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="forced tool use"): + config.map_tool_choice_values( + model="anthropic.claude-fable-5-1", tool_choice=tool_choice, drop_params=False + ) + + +@pytest.mark.parametrize("tool_choice", ["auto", "none"]) +def test_unforced_tool_choice_unaffected_on_fable_5_1_converse(local_model_cost_map, tool_choice): + config = AmazonConverseConfig() + + result = config.map_tool_choice_values( + model="anthropic.claude-fable-5-1", tool_choice=tool_choice, drop_params=True + ) + + assert result == ({"auto": {}} if tool_choice == "auto" else None) + + +@pytest.mark.parametrize( + "model", + ["anthropic.claude-fable-5-1", "us.anthropic.claude-fable-5-1"], +) +def test_response_format_avoids_native_and_forced_tool_choice_on_fable_5_1_converse( + local_model_cost_map, model +): + """Regression: Bedrock rejects both ``outputConfig`` structured output and forced + tool_choice for Fable 5.1, so response_format must map to a tool without a forced + tool_choice.""" + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + } + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "outputConfig" not in result + assert "tools" in result + assert "tool_choice" not in result + assert result.get("json_mode") is True + + +def test_forced_tool_choice_forwarded_on_converse_models_that_support_it( + local_model_cost_map, monkeypatch +): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + result = config.map_tool_choice_values( + model="anthropic.claude-fable-5", tool_choice="required", drop_params=False + ) + + assert result == {"any": {}} diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py index 7f91b49a6f5..639be272351 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -204,3 +204,69 @@ def test_should_forward_trusted_model_credentials_to_retrieve_provider_config(): assert response is mock_response litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"] assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials + + +@pytest.mark.asyncio +async def test_afile_content_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-download": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESDOWNLOADROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + class FakeS3Body: + def read(self): + return b'{"custom_id": "req-1"}' + + class FakeS3Client: + def get_object(self, Bucket, Key): + return {"Body": FakeS3Body()} + + def fake_boto3_client(service_name, **kwargs): + if service_name == "sts": + return FakeSTSClient() + return FakeS3Client() + + optional_params = { + "_litellm_internal_model_credentials": MappingProxyType({"s3_bucket_name": "safe-bucket"}), + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESDOWNLOADCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-download-role", + "aws_session_name": "litellm-files-download-session", + "aws_external_id": "external-id-files-download", + } + + with patch.object(boto3, "client", side_effect=fake_boto3_client) as mock_boto3_client: + response = await BedrockFilesHandler().afile_content( + file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"}, + optional_params=optional_params, + timeout=10.0, + max_retries=None, + ) + + s3_client_kwargs = next(call.kwargs for call in mock_boto3_client.call_args_list if call.args[0] == "s3") + assert s3_client_kwargs["aws_access_key_id"] == "ASIAFILESDOWNLOADROLE" + assert s3_client_kwargs["aws_session_token"] == "assumed-session-token" + assert response.content == b'{"custom_id": "req-1"}' diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index da13f265ee4..541c0db15d8 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2404,3 +2404,111 @@ class TestBedrockFilesS3SignatureEncoding: body=None, headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM], ) + + +def test_sign_s3_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 upload request.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-put": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESPUTROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESPUTCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-put-role", + "aws_session_name": "litellm-files-put-session", + "aws_external_id": "external-id-files-put", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, _signed_body = BedrockFilesConfig()._sign_s3_request( + content='{"custom_id": "req-1"}', + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIAFILESPUTROLE" in authorization + + +def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.files.transformation import ( + BedrockFilesConfig, + _BedrockS3RequestParams, + ) + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-files-get": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAFILESGETROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + request_params = _BedrockS3RequestParams.model_validate( + { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAFILESGETCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-get-role", + "aws_session_name": "litellm-files-get-session", + "aws_external_id": "external-id-files-get", + } + ) + assert request_params.aws_external_id == "external-id-files-get" + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers = BedrockFilesConfig()._sign_s3_get_request( + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + aws_region_name="us-east-1", + request_params=request_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIAFILESGETROLE" in authorization diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 8d07d38b1b6..09ebc1a3c95 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -935,7 +935,7 @@ def test_bedrock_messages_strips_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -970,7 +970,7 @@ def test_bedrock_messages_preserves_output_config_for_claude_4_6(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1003,7 +1003,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ) as mock_supports_factory: result = cfg.transform_anthropic_messages_request( @@ -1014,11 +1014,7 @@ def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): headers={}, ) - mock_supports_factory.assert_called_with( - model="us.anthropic.claude-opus-4-7", - custom_llm_provider="bedrock", - key="supports_output_config", - ) + mock_supports_factory.assert_called_with("us.anthropic.claude-opus-4-7", "supports_output_config") assert result["output_config"] == {"effort": "high"} @@ -1038,7 +1034,7 @@ def test_bedrock_messages_forwards_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1054,27 +1050,29 @@ def test_bedrock_messages_forwards_output_config(): def test_bedrock_messages_forwards_output_config_with_output_format(): - """``output_config`` is forwarded; ``output_format`` is converted to inline schema.""" + """Legacy ``output_format`` is forwarded as ``output_config.format`` on models + that support native structured outputs, alongside the effort key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + } optional_params = { "max_tokens": 4096, "output_config": {"effort": "low"}, - "output_format": { - "type": "json_schema", - "schema": { - "type": "object", - "properties": {"answer": {"type": "string"}}, - }, - }, + "output_format": schema_format, } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1085,12 +1083,14 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): headers={}, ) - assert result.get("output_config") == {"effort": "low"} + assert result.get("output_config") == {"effort": "low", "format": schema_format} assert "output_format" not in result + assert "answer" not in json.dumps(result["messages"]) def test_bedrock_messages_converts_output_config_format_to_inline_schema(): - """``output_config.format`` is consumed so Bedrock does not see an unknown nested key.""" + """Without native structured-output support, ``output_config.format`` falls back + to the inline schema so Bedrock does not see an unknown nested key.""" from unittest.mock import patch from litellm.types.router import GenericLiteLLMParams @@ -1110,8 +1110,8 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1146,7 +1146,7 @@ def test_bedrock_messages_normalizes_output_config_effort_for_opus( cfg = AmazonAnthropicClaudeMessagesConfig() with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1184,8 +1184,8 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", - return_value=True, + "litellm.llms.bedrock.common_utils._bedrock_model_supports", + side_effect=lambda _model, key: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1229,7 +1229,7 @@ def test_bedrock_messages_does_not_mutate_callers_output_config(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): cfg.transform_anthropic_messages_request( @@ -1271,7 +1271,7 @@ def test_bedrock_messages_strips_output_config_with_output_format(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=False, ): result = cfg.transform_anthropic_messages_request( @@ -1332,7 +1332,7 @@ def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): litellm.drop_params = True try: with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1375,7 +1375,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -1482,7 +1482,7 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): } with patch( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + "litellm.llms.bedrock.common_utils._bedrock_model_supports", return_value=True, ): result = cfg.transform_anthropic_messages_request( @@ -3104,3 +3104,149 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect(): break await asyncio.sleep(0.01) assert logging_obj.completion_start_time is not None + + +def test_bedrock_messages_forwards_output_config_format_natively(local_model_cost_map): + """Regression: on a model Bedrock enforces structured outputs for (Claude + Sonnet 4.5), ``output_config.format`` must be forwarded verbatim, not + silently rewritten into inline prompt text.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "zebra_count": {"type": "integer"}, + "is_tuesday": {"type": "boolean"}, + }, + "required": ["zebra_count", "is_tuesday"], + "additionalProperties": False, + }, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + assert "zebra_count" not in json.dumps(result["messages"]) + + +def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): + """Bedrock rejects ``output_config.format`` for the Claude 5 family, so the + schema falls back to the inline-text path instead of a deterministic 400.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema = { + "type": "object", + "properties": {"zebra_count": {"type": "integer"}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_config" not in result + last_content = result["messages"][-1]["content"] + assert json.loads(last_content[-1]["text"]) == schema + + +def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(local_model_cost_map): + """When a request carries both schema forms, the legacy top-level + ``output_format`` keeps winning, matching the pre-existing precedence.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + legacy_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"legacy_field": {"type": "string"}}}, + } + newer_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"newer_field": {"type": "string"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_format": legacy_format, + "output_config": {"format": newer_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": legacy_format} + assert "output_format" not in result + assert "newer_field" not in json.dumps(result) + + +def test_bedrock_messages_drop_params_keeps_native_output_config_format(local_model_cost_map, monkeypatch): + """``drop_params=True`` must not strip a natively forwarded + ``output_config.format`` on models without effort support (Sonnet 4.5).""" + import litellm + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setattr(litellm, "drop_params", True) + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": schema_format}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} + + +def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_model_cost_map): + """Sonnet 4.5 has native structured-output support but no effort support, so + a mixed ``output_config`` keeps ``format`` and drops ``effort``.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + schema_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + } + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 4096, + "output_config": {"format": schema_format, "effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert result.get("output_config") == {"format": schema_format} diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 7d07ac947b1..f854d806bdc 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest from botocore.auth import SigV4Auth from botocore.credentials import Credentials +from botocore.exceptions import NoCredentialsError import litellm from litellm.llms.bedrock.base_aws_llm import ( @@ -801,6 +802,23 @@ def test_get_request_headers_with_sigv4(): assert result == mock_request.prepare.return_value +def test_get_request_headers_without_credentials_or_bearer_token_raises_no_credentials(): + """Bearer-token auth needs no SigV4 principal, so `credentials` may be None. + Reaching the SigV4 branch with neither must fail the way botocore always + has instead of signing with a missing principal.""" + llm = BaseAWSLLM() + + with patch.dict(os.environ, {}, clear=True), pytest.raises(NoCredentialsError): + llm.get_request_headers( + credentials=None, + aws_region_name="us-west-2", + extra_headers=None, + endpoint_url="https://api.example.com", + data='{"prompt": "test"}', + headers={"Content-Type": "application/json"}, + ) + + def test_sigv4_matches_rust_golden_vector(): request = AWSRequest( method="POST", diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 389bf4a8e40..9302dc01abe 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -520,3 +520,97 @@ def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_stati assert merged["aws_secret_access_key"] == "caller-secret" assert merged["aws_session_token"] == "caller-token" assert merged["aws_region_name"] == "us-west-2" + + +def test_strip_unsupported_output_config_keeps_format_drops_effort(local_model_cost_map): + """On a model with neither effort flag, only the ``format`` key survives.""" + from litellm.llms.bedrock.common_utils import ( + strip_unsupported_bedrock_invoke_output_config_keys, + ) + + schema_format = {"type": "json_schema", "schema": {"type": "object"}} + body = {"output_config": {"effort": "high", "format": schema_format}} + + strip_unsupported_bedrock_invoke_output_config_keys( + model="anthropic.claude-3-haiku-20240307-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": schema_format} + + +def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_map): + """The legacy ``output_format`` wins over ``output_config.format`` when a + request carries both, matching the pre-existing precedence.""" + from litellm.llms.bedrock.common_utils import ( + apply_bedrock_invoke_structured_output, + ) + + legacy = {"type": "json_schema", "schema": {"type": "object", "properties": {"a": {"type": "string"}}}} + newer = {"type": "json_schema", "schema": {"type": "object", "properties": {"b": {"type": "string"}}}} + body = { + "messages": [{"role": "user", "content": "hi"}], + "output_format": legacy, + "output_config": {"format": newer}, + } + + apply_bedrock_invoke_structured_output( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + request_body=body, + ) + + assert body["output_config"] == {"format": legacy} + assert "output_format" not in body + + +def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied when signing batch API requests.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-batch-sign": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIABATCHSIGNROLE", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABATCHSIGNCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-batch-sign-role", + "aws_session_name": "litellm-batch-sign-session", + "aws_external_id": "external-id-batch-sign", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, signed_data = CommonBatchFilesUtils().sign_aws_request( + service_name="bedrock", + data={"jobName": "litellm-batch-job"}, + endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "ASIABATCHSIGNROLE" in authorization + assert signed_data == b'{"jobName": "litellm-batch-job"}' diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 4c92c52d556..7509e35e3f7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -1,7 +1,11 @@ import asyncio import concurrent.futures +import socket +import sys +from typing import Final import aiohttp +import aiohttp.abc import aiohttp.client_exceptions import aiohttp.http_exceptions import httpx @@ -1140,3 +1144,55 @@ async def test_stopped_loop_session_disposed_synchronously_on_recycle(): finally: await new_session.close() result["loop"].close() + + +class _CancellingResolver(aiohttp.abc.AbstractResolver): + """Cancels the given task (or, by default, aiohttp's shielded DNS child task) mid-lookup.""" + + def __init__(self, task_to_cancel: "asyncio.Task[object] | None" = None): + self._task_to_cancel: Final = task_to_cancel + + async def resolve( + self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET + ) -> list[aiohttp.abc.ResolveResult]: + target: Final = self._task_to_cancel or asyncio.current_task() + assert target is not None + target.cancel() + await asyncio.sleep(0) + raise OSError("resolver finished after the task was cancelled") + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.skipif( + sys.version_info < (3, 11), reason="Task.cancelling() is needed to tell the two cancellations apart" +) +async def test_internal_dns_cancellation_maps_to_connect_error(): + """A CancelledError the request task never asked for must surface as a mapped httpx transport error.""" + session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver())) + transport = LiteLLMAiohttpTransport(client=session) + try: + with pytest.raises(httpx.ConnectError): + await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/")) + current = asyncio.current_task() + assert current is not None and current.cancelling() == 0 + finally: + await transport.aclose() + + +@pytest.mark.asyncio +async def test_genuine_request_cancellation_still_propagates(): + """Cancelling the request task itself (client disconnect, shutdown) must still propagate unmapped.""" + current = asyncio.current_task() + assert current is not None + session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver(current))) + transport = LiteLLMAiohttpTransport(client=session) + try: + with pytest.raises(asyncio.CancelledError): + await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/")) + finally: + if sys.version_info >= (3, 11): + current.uncancel() + await transport.aclose() diff --git a/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py new file mode 100644 index 00000000000..064d9d58f0c --- /dev/null +++ b/tests/test_litellm/llms/dashscope/test_qwen_brand_aliases.py @@ -0,0 +1,331 @@ +import math + +import pytest + +import litellm +from litellm import completion, get_llm_provider +from litellm.llms.dashscope.chat.transformation import DashScopeChatConfig +from litellm.llms.dashscope.cost_calculator import ( + cost_per_token as dashscope_cost_per_token, +) +from litellm.llms.dashscope.embed.transformation import DashScopeEmbeddingConfig +from litellm.llms.dashscope.image_generation.transformation import ( + DashScopeImageGenerationConfig, +) +from litellm.llms.dashscope.qwen_ai_platform import ( + QWEN_AI_PLATFORM_API_BASE, + QWEN_AI_PLATFORM_IMAGE_API_BASE, + QWEN_AI_PLATFORM_RERANK_API_BASE, + QwenAIPlatformChatConfig, + QwenAIPlatformEmbeddingConfig, + QwenAIPlatformImageGenerationConfig, + QwenAIPlatformRerankConfig, +) +from litellm.llms.dashscope.qwencloud import ( + QWENCLOUD_API_BASE, + QWENCLOUD_IMAGE_API_BASE, + QWENCLOUD_RERANK_API_BASE, + QwenCloudChatConfig, + QwenCloudEmbeddingConfig, + QwenCloudImageGenerationConfig, + QwenCloudRerankConfig, +) +from litellm.llms.dashscope.rerank.transformation import DashScopeRerankConfig +from litellm.types.utils import LlmProviders, Usage +from litellm.utils import ProviderConfigManager + +DASHSCOPE_FAMILY_ENV_VARS = [ + "DASHSCOPE_API_KEY", + "DASHSCOPE_API_BASE", + "DASHSCOPE_API_BASE_RERANK", + "DASHSCOPE_API_BASE_IMAGE", + "QWENCLOUD_API_KEY", + "QWENCLOUD_API_BASE", + "QWENCLOUD_API_BASE_RERANK", + "QWENCLOUD_API_BASE_IMAGE", + "QWEN_AI_PLATFORM_API_KEY", + "QWEN_AI_PLATFORM_API_BASE", + "QWEN_AI_PLATFORM_API_BASE_RERANK", + "QWEN_AI_PLATFORM_API_BASE_IMAGE", +] + +BRAND_CASES = [ + pytest.param( + { + "provider": "qwencloud", + "enum": LlmProviders.QWENCLOUD, + "key_env": "QWENCLOUD_API_KEY", + "base_env": "QWENCLOUD_API_BASE", + "default_base": QWENCLOUD_API_BASE, + "default_rerank_base": QWENCLOUD_RERANK_API_BASE, + "default_image_base": QWENCLOUD_IMAGE_API_BASE, + "chat_config": QwenCloudChatConfig, + "embedding_config": QwenCloudEmbeddingConfig, + "rerank_config": QwenCloudRerankConfig, + "image_config": QwenCloudImageGenerationConfig, + }, + id="qwencloud", + ), + pytest.param( + { + "provider": "qwen_ai_platform", + "enum": LlmProviders.QWEN_AI_PLATFORM, + "key_env": "QWEN_AI_PLATFORM_API_KEY", + "base_env": "QWEN_AI_PLATFORM_API_BASE", + "default_base": QWEN_AI_PLATFORM_API_BASE, + "default_rerank_base": QWEN_AI_PLATFORM_RERANK_API_BASE, + "default_image_base": QWEN_AI_PLATFORM_IMAGE_API_BASE, + "chat_config": QwenAIPlatformChatConfig, + "embedding_config": QwenAIPlatformEmbeddingConfig, + "rerank_config": QwenAIPlatformRerankConfig, + "image_config": QwenAIPlatformImageGenerationConfig, + }, + id="qwen_ai_platform", + ), +] + + +@pytest.fixture(autouse=True) +def clear_dashscope_family_env(monkeypatch): + for env_var in DASHSCOPE_FAMILY_ENV_VARS: + monkeypatch.delenv(env_var, raising=False) + + +class TestQwenBrandProviderResolution: + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_get_llm_provider_resolves_brand_default_base(self, brand): + model, provider, api_key, api_base = get_llm_provider(f"{brand['provider']}/qwen-max", api_key="sk-explicit") + assert model == "qwen-max" + assert provider == brand["provider"] + assert api_key == "sk-explicit" + assert api_base == brand["default_base"] + + def test_dashscope_resolution_unchanged(self): + model, provider, api_key, api_base = get_llm_provider("dashscope/qwen-max", api_key="sk-explicit") + assert model == "qwen-max" + assert provider == "dashscope" + assert api_base == "https://dashscope.aliyuncs.com/compatible-mode/v1" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_brand_env_key_wins_over_dashscope_key(self, monkeypatch, brand): + monkeypatch.setenv(brand["key_env"], "sk-brand") + monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope") + _, _, api_key, _ = get_llm_provider(f"{brand['provider']}/qwen-max") + assert api_key == "sk-brand" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_dashscope_key_is_fallback(self, monkeypatch, brand): + monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope") + _, _, api_key, _ = get_llm_provider(f"{brand['provider']}/qwen-max") + assert api_key == "sk-dashscope" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_dashscope_api_base_does_not_leak_into_brand(self, monkeypatch, brand): + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://legacy.example.com/v1") + _, _, _, api_base = get_llm_provider(f"{brand['provider']}/qwen-max", api_key="sk-explicit") + assert api_base == brand["default_base"] + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_brand_api_base_env_wins(self, monkeypatch, brand): + monkeypatch.setenv(brand["base_env"], "https://brand.example.com/v1") + _, _, _, api_base = get_llm_provider(f"{brand['provider']}/qwen-max", api_key="sk-explicit") + assert api_base == "https://brand.example.com/v1" + + +class TestQwenBrandConfigDispatch: + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_chat_config(self, brand): + config = ProviderConfigManager.get_provider_chat_config("qwen-max", brand["enum"]) + assert isinstance(config, brand["chat_config"]) + assert isinstance(config, DashScopeChatConfig) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_embedding_config(self, brand): + config = ProviderConfigManager.get_provider_embedding_config(model="text-embedding-v3", provider=brand["enum"]) + assert isinstance(config, brand["embedding_config"]) + assert isinstance(config, DashScopeEmbeddingConfig) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_config(self, brand): + config = ProviderConfigManager.get_provider_rerank_config( + model="gte-rerank-v2", + provider=brand["enum"], + api_base=None, + present_version_params=[], + ) + assert isinstance(config, brand["rerank_config"]) + assert isinstance(config, DashScopeRerankConfig) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_config(self, brand): + config = ProviderConfigManager.get_provider_image_generation_config(model="qwen-image", provider=brand["enum"]) + assert isinstance(config, brand["image_config"]) + assert isinstance(config, DashScopeImageGenerationConfig) + + +class TestQwenBrandDefaultUrls: + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_chat_complete_url(self, brand): + url = brand["chat_config"]().get_complete_url( + api_base=None, + api_key="sk-test", + model="qwen-max", + optional_params={}, + litellm_params={}, + ) + assert url == f"{brand['default_base']}/chat/completions" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_embedding_complete_url(self, brand): + url = brand["embedding_config"]().get_complete_url( + api_base=None, + api_key="sk-test", + model="text-embedding-v3", + optional_params={}, + litellm_params={}, + ) + assert url == f"{brand['default_base']}/embeddings" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_embedding_ignores_dashscope_api_base(self, monkeypatch, brand): + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://legacy.example.com/v1") + url = brand["embedding_config"]().get_complete_url( + api_base=None, + api_key="sk-test", + model="text-embedding-v3", + optional_params={}, + litellm_params={}, + ) + assert url == f"{brand['default_base']}/embeddings" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_complete_url(self, brand): + url = brand["rerank_config"]().get_complete_url(api_base=None, model="gte-rerank-v2") + assert url == brand["default_rerank_base"] + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_rerank_env_override(self, monkeypatch, brand): + monkeypatch.setenv(f"{brand['base_env']}_RERANK", "https://rerank.example.com/v1/reranks") + url = brand["rerank_config"]().get_complete_url(api_base=None, model="gte-rerank-v2") + assert url == "https://rerank.example.com/v1/reranks" + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_complete_url(self, brand): + url = brand["image_config"]().get_complete_url( + api_base=None, + api_key="sk-test", + model="qwen-image", + optional_params={}, + litellm_params={}, + ) + assert url == brand["default_image_base"] + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_ignores_chat_compatible_api_base(self, brand): + url = brand["image_config"]().get_complete_url( + api_base=brand["default_base"], + api_key="sk-test", + model="qwen-image", + optional_params={}, + litellm_params={}, + ) + assert url == brand["default_image_base"] + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_validate_environment_requires_key(self, brand): + with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"): + brand["embedding_config"]().validate_environment( + headers={}, + model="text-embedding-v3", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + +class TestQwenBrandCostParity: + @pytest.fixture(autouse=True) + def setup_model_cost_map(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_get_model_info(self, brand): + model_info = litellm.get_model_info(f"{brand['provider']}/qwen-max") + dashscope_info = litellm.get_model_info("dashscope/qwen-max") + assert model_info["litellm_provider"] == brand["provider"] + assert model_info["input_cost_per_token"] == dashscope_info["input_cost_per_token"] + assert model_info["output_cost_per_token"] == dashscope_info["output_cost_per_token"] + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_flat_pricing_matches_dashscope(self, brand): + usage = Usage(prompt_tokens=1000, completion_tokens=500) + brand_costs = dashscope_cost_per_token(model="qwen-max", usage=usage, custom_llm_provider=brand["provider"]) + dashscope_costs = dashscope_cost_per_token(model="qwen-max", usage=usage) + assert brand_costs == dashscope_costs + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_tiered_pricing_matches_dashscope(self, brand): + usage = Usage(prompt_tokens=300000, completion_tokens=300000) + brand_costs = dashscope_cost_per_token(model="qwen-flash", usage=usage, custom_llm_provider=brand["provider"]) + dashscope_costs = dashscope_cost_per_token(model="qwen-flash", usage=usage) + assert brand_costs == dashscope_costs + tier_2 = litellm.get_model_info(f"{brand['provider']}/qwen-flash")["tiered_pricing"][1] + assert math.isclose(brand_costs[0], 300000 * tier_2["input_cost_per_token"], rel_tol=1e-10) + + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_public_cost_per_token_routes_to_dashscope_calculator(self, brand): + brand_costs = litellm.cost_per_token( + model=f"{brand['provider']}/qwen-max", + prompt_tokens=1000, + completion_tokens=500, + custom_llm_provider=brand["provider"], + ) + dashscope_costs = litellm.cost_per_token( + model="dashscope/qwen-max", + prompt_tokens=1000, + completion_tokens=500, + custom_llm_provider="dashscope", + ) + assert brand_costs == dashscope_costs + + +class TestQwenBrandCompletionMock: + @pytest.mark.respx() + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_completion_hits_brand_default_host(self, respx_mock, brand, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(f"{brand['default_base']}/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "qwen-turbo", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hey from LiteLLM!"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + }, + status_code=200, + ) + + response = completion( + model=f"{brand['provider']}/qwen-turbo", + messages=[{"role": "user", "content": "say hey from LiteLLM"}], + api_key="fake-brand-key", + ) + + assert response.choices[0].message.content == "Hey from LiteLLM!" + request = respx_mock.calls[0].request + assert request.url == f"{brand['default_base']}/chat/completions" + assert request.headers["Authorization"] == "Bearer fake-brand-key" diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 29ad8ee4b6e..e72642f7a04 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -62,6 +62,8 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-gemini-2-5-pro": ("22.321", "178.571", "22.321", "2.232"), "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), "databricks/databricks-kimi-k3": ("42.857", "214.286", "42.857", "4.286"), + "databricks/databricks-deepseek-v4-flash-0731": ("2.000", "4.000", "2.000", "0.400"), + "databricks/databricks-deepseek-v4-pro-0813": ("18.857", "56.571", "18.857", "1.886"), "databricks/databricks-glm-5-2": ("20.000", "62.857", "20.000", "3.714"), } PROMOTIONAL_DISCOUNT: Final = 0.80 diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 890df597933..a0e1616d4b2 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -24,6 +24,9 @@ OCR3_MODEL = "mistral/mistral-ocr-2512" OCR3_COST_PER_PAGE = 0.002 OCR3_ANNOTATION_COST_PER_PAGE = 0.003 +AZURE_DOC_AI_MODEL = "azure_ai/mistral-document-ai-2512" +AZURE_DOC_AI_COST_PER_PAGE = 0.003 + def _ocr_response(model: str, pages_processed: int) -> OCRResponse: return OCRResponse( @@ -33,6 +36,14 @@ def _ocr_response(model: str, pages_processed: int) -> OCRResponse: ) +def _annotated_ocr_response(model: str, pages_processed: int | None, annotation_pages: int) -> OCRResponse: + return OCRResponse( + pages=[], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, pages_processed_annotation=annotation_pages), + ) + + @pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) def test_model_info_ocr4_price(model: str) -> None: info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") @@ -79,3 +90,46 @@ def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) call_type="ocr", ) assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) + + +def test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 2, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(2 * OCR3_COST_PER_PAGE + 3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_only_response(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", 0, 3), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(3 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_ocr3_bills_annotation_pages_when_pages_processed_missing(local_model_cost_map) -> None: + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-ocr-2512", None, 4), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(4 * OCR3_ANNOTATION_COST_PER_PAGE) + + +def test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate(local_model_cost_map) -> None: + info = litellm.get_model_info(model=AZURE_DOC_AI_MODEL, custom_llm_provider="azure_ai") + assert info.get("annotation_cost_per_page") is None + assert info["ocr_cost_per_page"] == AZURE_DOC_AI_COST_PER_PAGE + cost = completion_cost( + completion_response=_annotated_ocr_response("mistral-document-ai-2512", 0, 1), + model=AZURE_DOC_AI_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + ) + assert cost == pytest.approx(AZURE_DOC_AI_COST_PER_PAGE) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..7dd6065063a 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1559,3 +1559,87 @@ class TestScanOnlyToolResults: assert data["messages"][3]["content"] == "page says [BLOCKED] here" assert data["messages"][3]["tool_call_id"] == "call_1" assert data["messages"][4]["content"] == "and then?" + + +class TestBuildBlockSseChunks: + """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" + + def _exc(self, original_response=None): + from litellm.exceptions import ModifyResponseException + + return ModifyResponseException( + message="Blocked by policy.", + model="gpt-5.4-mini", + request_data={}, + guardrail_name="test", + original_response=original_response, + ) + + def _payloads(self, chunks): + return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + + def test_standalone_block_uses_fresh_identity_and_zero_usage(self): + handler = OpenAIChatCompletionsHandler() + first, final = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False)) + assert first["id"].startswith("chatcmpl-") + assert first["model"] == "gpt-5.4-mini" + assert first["choices"][0]["delta"] == {"role": "assistant", "content": "Blocked by policy."} + assert first["choices"][0]["finish_reason"] is None + assert final["choices"][0]["delta"] == {} + assert final["choices"][0]["finish_reason"] == "content_filter" + assert final["usage"] == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + def test_continuation_reuses_stream_identity_and_real_usage(self): + handler = OpenAIChatCompletionsHandler() + yielded = [ + {"id": "chatcmpl-live", "created": 1724900000, "model": "gpt-5.4-mini-2026-01-01"}, + ] + original = yielded + [ + {"id": "chatcmpl-live", "usage": {"prompt_tokens": 11, "completion_tokens": 5}}, + ] + first, final = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=yielded + ) + ) + assert (first["id"], first["created"], first["model"]) == ( + "chatcmpl-live", + 1724900000, + "gpt-5.4-mini-2026-01-01", + ) + assert first["choices"][0]["delta"] == {"content": "Blocked by policy."} + assert final["id"] == "chatcmpl-live" + assert final["usage"] == {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16} + + +class TestCheckStreamingHasEnded: + """_check_streaming_has_ended lets end_of_stream_only withhold the finish chunk until moderation""" + + def test_empty_and_content_only_chunks_are_not_ended(self): + handler = OpenAIChatCompletionsHandler() + assert handler._check_streaming_has_ended([]) is False + content_only = [ + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]}, + {"id": "chatcmpl-live", "choices": []}, + {"id": "chatcmpl-live", "usage": {"prompt_tokens": 1, "completion_tokens": 1}}, + ] + assert handler._check_streaming_has_ended(content_only) is False + + def test_dict_finish_chunk_marks_stream_ended(self): + handler = OpenAIChatCompletionsHandler() + chunks = [ + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]}, + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ] + assert handler._check_streaming_has_ended(chunks) is True + + def test_object_finish_chunk_marks_stream_ended(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + chunks = [ + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")] + ) + ] + assert handler._check_streaming_has_ended(chunks) is True diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 3f346b5e8e7..9737d63cc26 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -145,6 +145,69 @@ class TestGetOptionalParamsIntegration: assert regular_params.get("user") == "my-end-user" assert responses_params.get("user") == "my-end-user" + def test_reasoning_effort_supported_for_unknown_model_alias(self): + """An openai/-routed model litellm doesn't recognize is likely a proxy alias: + reasoning_effort must be forwarded so the server decides support.""" + from litellm.llms.openai.openai import OpenAIConfig + + supported_params = OpenAIConfig().get_supported_openai_params( + "my-claude-alias" + ) + assert "reasoning_effort" in supported_params + + def test_reasoning_effort_not_supported_for_known_non_reasoning_models(self): + """Known OpenAI models keep failing closed client-side.""" + from litellm.llms.openai.openai import OpenAIConfig + + config = OpenAIConfig() + assert "reasoning_effort" not in config.get_supported_openai_params("gpt-4o") + assert "reasoning_effort" not in config.get_supported_openai_params( + "responses/gpt-4.1-mini" + ) + + def test_reasoning_effort_not_inherited_by_openai_compatible_subclasses(self): + """Providers subclassing either openai config keep their own reasoning_effort gating + for their models, which are all unknown to the openai catalog.""" + from litellm.llms.openai.openai import OpenAIConfig + + class InheritingDispatcherConfig(OpenAIConfig): + pass + + class InheritingGPTConfig(OpenAIGPTConfig): + pass + + assert "reasoning_effort" not in InheritingDispatcherConfig().get_supported_openai_params( + "some-unknown-model" + ) + assert "reasoning_effort" not in InheritingGPTConfig().get_supported_openai_params( + "some-unknown-model" + ) + + def test_reasoning_effort_forwarded_in_optional_params_for_unknown_model_alias( + self, + ): + """Regression test for reasoning_effort raising UnsupportedParamsError + client-side for openai/-prefixed proxy aliases before any HTTP request.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="my-claude-alias", + custom_llm_provider="openai", + reasoning_effort="low", + ) + assert optional_params.get("reasoning_effort") == "low" + + def test_reasoning_effort_still_rejected_for_known_non_reasoning_model(self): + """A real OpenAI model that doesn't reason still rejects the param client-side.""" + from litellm.utils import get_optional_params + + with pytest.raises(litellm.utils.UnsupportedParamsError): + get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + reasoning_effort="low", + ) + class TestOpenAIChatCompletionStreamingHandler: """Tests for OpenAIChatCompletionStreamingHandler.chunk_parser()""" @@ -808,6 +871,134 @@ class TestCacheControlPreservationForCustomEndpoint: assert all("cache_control" not in m for m in body["messages"]) +class TestToolChoiceWithoutToolsDropped: + def setup_method(self): + self.config = OpenAIGPTConfig() + + @staticmethod + def _pi_compact_summarization_messages(): + return [ + { + "role": "system", + "content": "You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.", + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "\n[User]: Reply with exactly: ok-1\n\n[Assistant]: ok-1\n\n\nThe messages above are a conversation to summarize.", + } + ], + }, + ] + + def _transform(self, optional_params, config=None, model="gpt-5.6-sol"): + return (config or self.config).transform_request( + model=model, + messages=self._pi_compact_summarization_messages(), + optional_params=optional_params, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + + def test_pi_compact_shape_drops_tool_choice_none_without_tools(self): + body = self._transform( + { + "stream": True, + "stream_options": {"include_usage": True}, + "store": False, + "max_completion_tokens": 13107, + "tool_choice": "none", + } + ) + assert "tool_choice" not in body + assert "tools" not in body + assert body["model"] == "gpt-5.6-sol" + assert body["stream"] is True + assert body["stream_options"] == {"include_usage": True} + assert body["store"] is False + assert body["max_completion_tokens"] == 13107 + + def test_drops_tool_choice_auto_without_tools(self): + body = self._transform({"tool_choice": "auto"}) + assert "tool_choice" not in body + + def test_drops_named_function_tool_choice_without_tools(self): + body = self._transform( + {"tool_choice": {"type": "function", "function": {"name": "get_weather"}}} + ) + assert "tool_choice" not in body + + def test_drops_tool_choice_but_keeps_empty_tools_array(self): + body = self._transform({"tools": [], "tool_choice": "none"}) + assert "tool_choice" not in body + assert body["tools"] == [] + + def test_gpt5_config_drops_tool_choice_without_tools(self): + body = self._transform({"tool_choice": "none"}, config=OpenAIGPT5Config()) + assert "tool_choice" not in body + + @pytest.mark.parametrize( + "tool_choice", + [ + "none", + "auto", + "required", + {"type": "function", "function": {"name": "get_weather"}}, + ], + ) + def test_preserves_tool_choice_when_tools_present(self, tool_choice): + tools = [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {}}, + } + ] + body = self._transform({"tools": tools, "tool_choice": tool_choice}) + assert body["tool_choice"] == tool_choice + assert body["tools"] == tools + + def test_preserves_tool_choice_with_legacy_functions(self): + functions = [{"name": "get_weather", "parameters": {}}] + body = self._transform({"functions": functions, "tool_choice": "auto"}) + assert body["tool_choice"] == "auto" + assert body["functions"] == functions + + def test_preserves_function_call_without_functions(self): + body = self._transform({"function_call": "none"}) + assert body["function_call"] == "none" + + @pytest.mark.asyncio + async def test_async_transform_drops_tool_choice_without_tools(self): + body = await self.config.async_transform_request( + model="gpt-5.6-sol", + messages=self._pi_compact_summarization_messages(), + optional_params={"stream": True, "tool_choice": "none"}, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + assert "tool_choice" not in body + + @pytest.mark.asyncio + async def test_async_transform_preserves_tool_choice_when_tools_present(self): + tools = [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {}}, + } + ] + body = await self.config.async_transform_request( + model="gpt-5.6-sol", + messages=self._pi_compact_summarization_messages(), + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + assert body["tool_choice"] == "auto" + assert body["tools"] == tools + + class TestToolMessageImageHoisting: """transform_request moves tool-message images into a following user message (OpenAI-compatible APIs only accept text in role:"tool" messages).""" @@ -975,3 +1166,118 @@ class TestOpenAIPromptCacheBreakpointChatPath: assert request["messages"][1]["content"] == [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}] assert request["extra_body"] == {"prompt_cache_options": self.EXPLICIT} assert "prompt_cache_options" not in request + + +class TestToolSchemaCombinatorFlatteningForOpenAI: + """ + Regression tests for LIT-6488: OpenAI's chat completions validator rejects + tool parameters carrying a top-level anyOf/oneOf/allOf for every model + family, GPT-5 included, unlike the Responses API. + """ + + def setup_method(self): + self.config = OpenAIGPTConfig() + + @pytest.fixture(autouse=True) + def _clean_openai_base_env(self, monkeypatch): + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None, raising=False) + + @staticmethod + def _anyof_tool(): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def _transform(self, config, model, litellm_params, tools): + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools}, + litellm_params=litellm_params, + headers={}, + ) + + def test_flattens_top_level_anyof_for_hosted_openai(self): + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [self._anyof_tool()] + ) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert request["tools"][0]["function"]["name"] == "automation_update" + + def test_gpt5_family_flattens_on_chat_completions(self): + request = self._transform( + OpenAIGPT5Config(), "gpt-5.6", {"custom_llm_provider": "openai", "api_base": None}, [self._anyof_tool()] + ) + assert "anyOf" not in request["tools"][0]["function"]["parameters"] + + def test_custom_api_base_keeps_union(self): + tool = self._anyof_tool() + request = self._transform( + self.config, + "gpt-4o", + {"custom_llm_provider": "openai", "api_base": "http://localhost:8000/v1"}, + [tool], + ) + assert request["tools"][0]["function"]["parameters"] == self._anyof_tool()["function"]["parameters"] + + def test_non_openai_provider_keeps_union(self): + request = self._transform( + self.config, "some-oss-model", {"custom_llm_provider": "groq", "api_base": None}, [self._anyof_tool()] + ) + assert request["tools"][0]["function"]["parameters"] == self._anyof_tool()["function"]["parameters"] + + def test_caller_tool_dict_is_not_mutated(self): + tool = self._anyof_tool() + self._transform(self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool]) + assert tool == self._anyof_tool() + + def test_clean_object_schema_passes_through_as_same_object(self): + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + request = self._transform( + self.config, "gpt-4o", {"custom_llm_provider": "openai", "api_base": None}, [tool] + ) + assert request["tools"][0] is tool + + @pytest.mark.asyncio + async def test_async_transform_request_flattens_for_hosted_openai(self): + request = await self.config.async_transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": [self._anyof_tool()]}, + litellm_params={"custom_llm_provider": "openai", "api_base": None}, + headers={}, + ) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..315b6948bd8 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -828,6 +828,24 @@ class MockPassThroughGuardrail(CustomGuardrail): return inputs +class MockRecordingGuardrail(MockPassThroughGuardrail): + """Pass-through guardrail that records every apply_guardrail inputs payload""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.seen_inputs: List[GenericGuardrailAPIInputs] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_inputs.append(inputs) + return inputs + + class TestOpenAIResponsesHandlerStreamingOutputProcessing: """Test streaming output processing functionality""" @@ -1104,6 +1122,80 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: output_text = result[-1]["response"]["output"][0]["content"][0]["text"] assert output_text == original_text + @pytest.mark.asyncio + async def test_failed_stream_scans_delta_text(self): + """A stream ending in response.failed has text only in delta events; the + fallback scan must assemble and scan it instead of skipping on an empty string.""" + handler = OpenAIResponsesHandler() + guardrail = MockRecordingGuardrail(guardrail_name="test") + + responses_so_far = [ + {"type": "response.created", "response": {"id": "resp_123"}}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_123"}}, + { + "type": "response.output_text.delta", + "item_id": "msg_123", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_123", + "output_index": 0, + "content_index": 0, + "delta": " world", + }, + {"type": "response.failed", "response": {"id": "resp_123", "status": "failed"}}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result == responses_so_far + assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["Hello world"]] + + def test_get_streaming_string_so_far_prefers_done_text_over_deltas(self): + """The done event repeats the whole part, so deltas must not be double counted; + a part with no done event yet still contributes its joined deltas.""" + handler = OpenAIResponsesHandler() + + events = [ + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": " world", + }, + { + "type": "response.output_text.done", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": "Hello world", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_2", + "output_index": 1, + "content_index": 0, + "delta": "; unfinished", + }, + ] + + assert handler.get_streaming_string_so_far(events) == "Hello world; unfinished" + class TestGetStructuredMessages: """Test the get_structured_messages method for Responses API handler.""" @@ -1229,3 +1321,219 @@ class TestOpenAIResponsesHandlerToolInjection: names = [t.get("name") for t in result["tools"]] assert "get_weather" in names assert "injected_tool" in names + + +class TestBuildBlockSseChunks: + """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" + + def _exc(self, original_response=None): + from litellm.exceptions import ModifyResponseException + + return ModifyResponseException( + message="Blocked by policy.", + model="gpt-5.4-mini", + request_data={}, + guardrail_name="test", + original_response=original_response, + ) + + def _payloads(self, chunks): + import json + + return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + + def test_standalone_block_emits_complete_synthetic_stream(self): + handler = OpenAIResponsesHandler() + payloads = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False)) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.created" + assert types[-1] == "response.completed" + completed = payloads[-1]["response"] + assert completed["id"].startswith("resp_") + assert completed["model"] == "gpt-5.4-mini" + assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." + + def test_continuation_appends_item_at_next_output_index_with_real_usage(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini-2026-01-01"}}, + {"type": "response.output_item.added", "output_index": 2, "item": {"id": "msg_orig"}}, + ] + original = yielded + [ + { + "type": "response.completed", + "response": { + "id": "resp_live", + "model": "gpt-5.4-mini-2026-01-01", + "output": [], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + }, + } + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert "response.created" not in types + assert types[0] == "response.output_item.done" + assert payloads[0]["output_index"] == 2 + assert payloads[0]["item"]["id"] == "msg_orig" + assert payloads[0]["item"]["status"] == "completed" + assert types[1] == "response.output_item.added" + assert payloads[1]["output_index"] == 3 + completed = payloads[-1]["response"] + assert completed["id"] == "resp_live" + assert completed["model"] == "gpt-5.4-mini-2026-01-01" + assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + + def test_continuation_reads_usage_from_typed_completed_event(self): + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + handler = OpenAIResponsesHandler() + original = [ + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse.model_validate( + { + "id": "resp_live", + "created_at": 1, + "model": "gpt-5.4-mini", + "output": [], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + } + ), + ) + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=[] + ) + ) + completed = payloads[-1]["response"] + assert completed["usage"]["input_tokens"] == 7 + assert completed["usage"]["output_tokens"] == 21 + assert completed["usage"]["total_tokens"] == 28 + + def test_continuation_closes_open_item_given_pydantic_events_with_enum_types(self): + from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ContentPartAddedEvent, + OutputItemAddedEvent, + OutputTextDeltaEvent, + ResponsesAPIStreamEvents, + ) + + handler = OpenAIResponsesHandler() + open_item = GenericResponseOutputItem.model_validate( + {"type": "message", "id": "msg_live", "status": "in_progress", "role": "assistant", "content": []} + ) + yielded = [ + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=open_item + ), + ContentPartAddedEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id="msg_live", + output_index=0, + content_index=0, + part=BaseLiteLLMOpenAIResponseObject.model_validate( + {"type": "output_text", "text": "", "annotations": []} + ), + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_live", + output_index=0, + content_index=0, + delta="partial ", + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_live", + output_index=0, + content_index=0, + delta="text", + ), + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[:3] == [ + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + ] + assert payloads[0]["text"] == "partial text" + assert payloads[2]["item"]["id"] == "msg_live" + assert payloads[2]["item"]["status"] == "completed" + assert payloads[2]["item"]["content"][0]["text"] == "partial text" + assert types[3] == "response.output_item.added" + assert payloads[3]["output_index"] == 1 + + def test_continuation_closes_open_function_call_as_incomplete(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}}, + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "fc_live", + "type": "function_call", + "status": "in_progress", + "call_id": "call_1", + "name": "run_payment", + "arguments": "", + }, + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_live", + "output_index": 0, + "delta": '{"amount": 100}', + }, + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.output_item.done" + closed = payloads[0]["item"] + assert closed["id"] == "fc_live" + assert closed["type"] == "function_call" + assert closed["status"] == "incomplete" + assert closed["name"] == "run_payment" + assert "content" not in closed + assert types[1] == "response.output_item.added" + assert payloads[1]["output_index"] == 1 + assert types[-1] == "response.completed" + + def test_continuation_without_open_item_emits_no_closing_events(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}}, + {"type": "response.in_progress", "response": {"id": "resp_live"}}, + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.output_item.added" + assert types[-1] == "response.completed" + dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] + assert len(dones) == 1 + assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index 1a90db7c1fe..b0ffd1845fe 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -220,6 +220,111 @@ class TestOpenAIResponsesAPIConfig: assert result["input"] == input_clean + def test_transform_drops_foreign_tool_call_item_ids(self): + """Replayed tool call items whose ids are not OpenAI-shaped (e.g. + Anthropic toolu_/srvtoolu_ ids after a router fallback) must be sent + without an id: OpenAI 400s foreign ids ("Expected an ID that begins + with 'fc'") but accepts the items with no id at all. Genuine fc_/ctc_ + ids and non-tool-call items pass through untouched.""" + replayed_input = [ + {"role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + { + "type": "function_call", + "id": "toolu_01Foreign", + "call_id": "toolu_01Foreign", + "name": "get_weather", + "arguments": '{"city": "SF"}', + }, + {"type": "function_call_output", "call_id": "toolu_01Foreign", "output": "sunny"}, + { + "type": "custom_tool_call", + "id": "srvtoolu_01Foreign", + "call_id": "srvtoolu_01Foreign", + "name": "apply_patch", + "input": "patch", + }, + { + "type": "function_call", + "id": "fc_genuine", + "call_id": "call_genuine", + "name": "get_weather", + "arguments": "{}", + }, + {"type": "message", "id": "msg_1", "role": "assistant", "content": []}, + ] + + result = self.config.transform_responses_api_request( + model=self.model, + input=replayed_input, + response_api_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + assert "id" not in result["input"][1] + assert result["input"][1]["call_id"] == "toolu_01Foreign" + assert "id" not in result["input"][3] + assert result["input"][3]["call_id"] == "srvtoolu_01Foreign" + assert result["input"][4]["id"] == "fc_genuine" + assert result["input"][5]["id"] == "msg_1" + assert replayed_input[1]["id"] == "toolu_01Foreign" + assert replayed_input[3]["id"] == "srvtoolu_01Foreign" + + def test_transform_keeps_foreign_tool_call_item_ids_for_other_providers(self): + """Providers reusing this config that do not enforce OpenAI's id + shapes must keep replayed ids untouched.""" + from litellm.types.utils import LlmProviders + + class _OpenRouterLikeConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.OPENROUTER + + replayed_input = [ + { + "type": "function_call", + "id": "toolu_01Foreign", + "call_id": "toolu_01Foreign", + "name": "get_weather", + "arguments": "{}", + } + ] + + result = _OpenRouterLikeConfig().transform_responses_api_request( + model="openrouter/some-model", + input=replayed_input, + response_api_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + assert result["input"][0]["id"] == "toolu_01Foreign" + + def test_transform_compact_drops_foreign_tool_call_item_ids(self): + """The compact request path replays input the same way, so it must + apply the same id drop.""" + replayed_input = [ + { + "type": "function_call", + "id": "toolu_01Foreign", + "call_id": "toolu_01Foreign", + "name": "get_weather", + "arguments": "{}", + } + ] + + _url, data = self.config.transform_compact_response_api_request( + model=self.model, + input=replayed_input, + response_api_optional_request_params={}, + api_base="https://api.openai.com/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "id" not in data["input"][0] + assert data["input"][0]["call_id"] == "toolu_01Foreign" + def test_transform_streaming_response(self): """Test streaming response transformation""" # Test with a text delta event diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index 8a9ae4dae6d..62b4d003b45 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -2,6 +2,7 @@ Tests for Parallel AI Search API integration (v1 endpoint). """ +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -30,13 +31,41 @@ MOCK_V1_RESPONSE = { } -def _mock_response(): +def _mock_response(payload=None): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.json.return_value = MOCK_V1_RESPONSE + mock_response.json.return_value = payload if payload is not None else MOCK_V1_RESPONSE return mock_response +@pytest.fixture +def httpx_transport(monkeypatch): + monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary. + litellm, + "disable_aiohttp_transport", + True, + ) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.fixture +def bundled_cost_map(monkeypatch): + """Price lookups against the bundled cost map. + + litellm caches model-info lookups, so swapping ``model_cost`` only takes + effect once those caches are invalidated -- on the way in and back out. + """ + from litellm.utils import _invalidate_model_cost_lowercase_map + + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + yield + monkeypatch.undo() + _invalidate_model_cost_lowercase_map() + + class TestParallelAISearch: @pytest.fixture(autouse=True) def _set_api_key(self, monkeypatch): @@ -135,9 +164,7 @@ class TestParallelAISearch: json_data = mock_post.call_args.kwargs.get("json") assert json_data["mode"] == "basic" - @pytest.mark.parametrize( - "processor,expected_mode", [("base", "basic"), ("pro", "advanced")] - ) + @pytest.mark.parametrize("processor,expected_mode", [("base", "basic"), ("pro", "advanced")]) @pytest.mark.asyncio async def test_legacy_processor_maps_to_mode(self, processor, expected_mode): with patch( @@ -222,9 +249,7 @@ class TestParallelAISearch: "arxiv.org", "nature.com", ] - assert advanced_settings["source_policy"]["exclude_domains"] == [ - "reddit.com" - ] + assert advanced_settings["source_policy"]["exclude_domains"] == ["reddit.com"] assert advanced_settings["excerpt_settings"]["max_chars_per_result"] == 1500 assert "max_results" not in json_data @@ -306,10 +331,7 @@ class TestParallelAISearch: ) call_args = mock_post.call_args - assert ( - call_args.kwargs["url"] - == "https://proxy.internal.example.com/v1/search" - ) + assert call_args.kwargs["url"] == "https://proxy.internal.example.com/v1/search" @pytest.mark.asyncio async def test_caller_api_base_without_key_is_refused(self, monkeypatch): @@ -338,3 +360,147 @@ class TestParallelAISearch: query="AI developments", search_provider="parallel_ai", ) + + @pytest.mark.asyncio + async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self, respx_mock, httpx_transport): + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) + + await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + objective="find peer-reviewed AI research", + include_domains=["arxiv.org"], + after_date="2026-01-01", + location="gb", + fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True}, + client_model="claude-fable-5", + ) + + json_data = json.loads(route.calls[0].request.content) + assert json_data["objective"] == "find peer-reviewed AI research" + assert json_data["client_model"] == "claude-fable-5" + + advanced_settings = json_data["advanced_settings"] + assert advanced_settings["location"] == "gb" + assert advanced_settings["fetch_policy"] == { + "max_age_seconds": 600, + "disable_cache_fallback": True, + } + assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"] + assert advanced_settings["source_policy"]["after_date"] == "2026-01-01" + + assert "include_domains" not in json_data + assert "after_date" not in json_data + assert "location" not in json_data + assert "fetch_policy" not in json_data + + @pytest.mark.asyncio + async def test_response_preserves_raw_parallel_fields(self, respx_mock, httpx_transport): + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + ) + + dumped = response.model_dump() + assert dumped["search_id"] == "search_abc123" + assert dumped["session_id"] == "session_xyz" + assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}] + + first = response.results[0].model_dump() + assert first["excerpts"] == ["First excerpt.", "Second excerpt."] + + @pytest.mark.asyncio + async def test_response_normalizes_null_result_fields(self, respx_mock, httpx_transport): + response_payload = { + **MOCK_V1_RESPONSE, + "results": [{"url": None, "title": None, "publish_date": None, "excerpts": None}], + } + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + ) + + assert len(response.results) == 1 + result = response.results[0] + assert result.url == "" + assert result.title == "" + assert result.snippet == "" + assert result.date is None + assert result.model_dump()["excerpts"] == () + + @pytest.mark.parametrize( + "mode,usage,max_results,expected_cost", + [ + ("turbo", [{"name": "sku_search", "count": 1}], None, 0.001), + ("fast", [{"name": "sku_search", "count": 1}], None, 0.001), + ("basic", [{"name": "sku_search", "count": 1}], None, 0.005), + ("advanced", [{"name": "sku_search", "count": 1}], None, 0.005), + ( + "basic", + [ + {"name": "sku_search", "count": 1}, + {"name": "sku_search_additional_results", "count": 2}, + ], + 20, + 0.007, + ), + ("basic", None, 20, 0.015), + ], + ) + @pytest.mark.asyncio + async def test_search_cost_uses_mode_and_provider_usage( + self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport + ): + response_payload = {**MOCK_V1_RESPONSE, "usage": usage} + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode=mode, + max_results=max_results, + ) + + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + @pytest.mark.asyncio + async def test_search_cost_treats_keyword_queries_as_one_request( + self, bundled_cost_map, respx_mock, httpx_transport + ): + response_payload = { + **MOCK_V1_RESPONSE, + "usage": [{"name": "sku_search", "count": 1}], + } + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query=["AI developments", "machine learning trends"], + search_provider="parallel_ai", + mode="basic", + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.005) + + @pytest.mark.asyncio + async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport): + """`_parallel_ai_usage` prices the request, so a caller must not be able to set it. + + The provider reports no usage here, which is the case where a caller-supplied + value would otherwise survive into the cost calculation. + """ + response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"} + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode="basic", + _parallel_ai_usage=[{"name": "sku_search", "count": 0}], + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.005) + assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content) diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py new file mode 100644 index 00000000000..72c69fc622c --- /dev/null +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py @@ -0,0 +1,191 @@ +"""Gateway coverage for Parallel AI Search.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Final +from unittest.mock import AsyncMock + +import httpx +import pytest +from fastapi.testclient import TestClient + +import litellm +from litellm import Router +from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.utils import LlmProviders + +PARALLEL_SEARCH_URL: Final = "https://api.parallel.ai/v1/search" + + +@pytest.fixture +def client() -> TestClient: + return TestClient(proxy_server.app, raise_server_exceptions=False) + + +@pytest.fixture +def auth_as() -> Iterator[None]: + async def _authorized_request() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-sk-test", + user_id="parallel-test-user", + ) + + previous: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth) + proxy_server.app.dependency_overrides[user_api_key_auth] = _authorized_request + try: + yield + finally: + if previous is None: + proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) + else: + proxy_server.app.dependency_overrides[user_api_key_auth] = previous + + +def _parallel_search_body() -> dict[str, object]: + return { + "search_id": "search_parallel_gateway", + "results": [ + { + "url": "https://example.com/parallel", + "title": "Parallel result", + "publish_date": "2026-08-13", + "excerpts": ["First excerpt", "Second excerpt"], + } + ], + "usage": [{"name": "sku_search", "count": 1}], + } + + +def _parallel_router(mode: str = "turbo") -> Router: + return Router( + model_list=[], + search_tools=[ + { + "search_tool_name": "parallel-search", + "litellm_params": { + "search_provider": "parallel_ai", + "api_key": "parallel-search-key", + "mode": mode, + }, + } + ], + num_retries=0, + ) + + +def _mock_async_post( + monkeypatch, + *, + url: str, + response_body: dict[str, object], +) -> AsyncMock: + response = httpx.Response( + status_code=200, + json=response_body, + request=httpx.Request("POST", url), + ) + mock_post = AsyncMock(return_value=response) + monkeypatch.setattr(AsyncHTTPHandler, "post", mock_post) + return mock_post + + +def test_parallel_search_gateway_route(client, auth_as, monkeypatch): + """The named search route selects its configured Parallel Search tool. + + The tool-level `mode` must survive the router hop, so the upstream request + is sent as `turbo` rather than falling back to the adapter default. + """ + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + monkeypatch.setattr(proxy_server, "llm_router", _parallel_router()) + mock_post = _mock_async_post( + monkeypatch, + url=PARALLEL_SEARCH_URL, + response_body=_parallel_search_body(), + ) + + response = client.post( + "/v1/search/parallel-search", + json={"query": "Parallel AI news", "max_results": 3}, + ) + + assert response.status_code == 200, response.text + assert response.json()["results"] == [ + { + "title": "Parallel result", + "url": "https://example.com/parallel", + "snippet": "First excerpt ... Second excerpt", + "date": "2026-08-13", + "last_updated": None, + "excerpts": ["First excerpt", "Second excerpt"], + } + ] + + request_kwargs = mock_post.await_args.kwargs + assert request_kwargs["url"] == PARALLEL_SEARCH_URL + assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key" + assert request_kwargs["json"] == { + "objective": "Parallel AI news", + "search_queries": ["Parallel AI news"], + "mode": "turbo", + "advanced_settings": {"max_results": 3}, + } + + +@pytest.mark.asyncio +async def test_web_search_interception_executes_parallel_search(monkeypatch): + """An intercepted web-search call uses the configured Parallel Search tool.""" + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + monkeypatch.setattr(proxy_server, "llm_router", _parallel_router(mode="fast")) + mock_post = _mock_async_post( + monkeypatch, + url=PARALLEL_SEARCH_URL, + response_body=_parallel_search_body(), + ) + logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI], + search_tool_name="parallel-search", + ) + + plan = await logger.async_build_responses_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "fc_parallel", + "call_id": "fc_parallel", + "type": "function_call", + "name": "litellm_web_search", + "arguments": '{"query":"Parallel AI news"}', + "input": {"query": "Parallel AI news"}, + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "Research Parallel"}], + response=None, + optional_params={"tools": [{"type": "function", "name": "litellm_web_search"}]}, + logging_obj=None, + stream=False, + kwargs={"custom_llm_provider": "openai"}, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.messages[-1] == { + "type": "function_call_output", + "call_id": "fc_parallel", + "output": ( + "Title: Parallel result\nURL: https://example.com/parallel\nSnippet: First excerpt ... Second excerpt" + ), + } + + request_kwargs = mock_post.await_args.kwargs + assert request_kwargs["url"] == PARALLEL_SEARCH_URL + assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key" + assert request_kwargs["json"]["mode"] == "fast" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 8c1de12e7d9..4679b978f78 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1096,10 +1096,13 @@ def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): "gemini-3.5-flash", "gemini-3.6-flash", "gemini-3.7-flash", + "gemini-3.8-flash", "vertex_ai/gemini-3.5-flash", "vertex_ai/gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", "gemini/gemini-3.5-flash", "gemini/gemini-3.7-flash", + "gemini/gemini-3.8-flash", ], ) def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index bd07bec900f..d2788408e09 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1185,6 +1185,18 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): } +def test_vertex_ai_map_thinking_param_without_budget_tokens_for_gemini_3(): + v = VertexGeminiConfig() + result = v.map_openai_params( + non_default_params={"thinking": {"type": "enabled"}}, + optional_params={}, + model="gemini-3.5-flash", + drop_params=False, + ) + + assert result["thinkingConfig"] == {"includeThoughts": True} + + def test_vertex_ai_map_tools(): v = VertexGeminiConfig() optional_params = {} diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index cc923f05831..d1d751989ea 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -195,6 +195,22 @@ def test_set_schema_property_ordering_with_excessive_nesting(): set_schema_property_ordering(schema) +def test_set_schema_property_ordering_skips_non_dict_property_values(): + """Non-dict property values must be skipped, not recursed into (they used to raise).""" + schema = { + "properties": { + "a": "hello", + "b": {"type": "string"}, + "c": ["x"], + "d": "a string mentioning items", + } + } + + result = set_schema_property_ordering(schema) + + assert result["propertyOrdering"] == ["a", "b", "c", "d"] + + def test_build_vertex_schema(): """Test build_vertex_schema with a sample schema""" from litellm.llms.vertex_ai.common_utils import _build_vertex_schema diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 05da22a73fd..fba337b5f2c 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -1,3 +1,4 @@ +import base64 from unittest.mock import MagicMock, Mock, patch import httpx @@ -126,6 +127,48 @@ class TestVertexAITextToSpeechConfig: assert voice_dict == voice_input +@pytest.mark.parametrize( + ("audio", "expected_content_type"), + [ + (b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00", "audio/wav"), + (b"\xff\xfb\x90\x64\x00\x00\x00\x00", "audio/mpeg"), + (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"), + (b"fLaC\x00\x00\x00\x22", "audio/flac"), + ], +) +def test_transform_text_to_speech_response_labels_content_type(audio, expected_content_type): + raw_response = httpx.Response( + status_code=200, + json={"audioContent": base64.b64encode(audio).decode()}, + ) + + result = VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert result.response.headers["content-type"] == expected_content_type + assert result.response.content == audio + + +def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): + raw_pcm = b"\x00\x01\x02\x03\x04\x05\x06\x07" + raw_response = httpx.Response( + status_code=200, + json={"audioContent": base64.b64encode(raw_pcm).decode()}, + ) + + result = VertexAITextToSpeechConfig().transform_text_to_speech_response( + model="vertex_ai/chirp", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert "content-type" not in result.response.headers + assert result.response.content == raw_pcm + + @patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 552ca98441f..9419f88a981 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -727,3 +727,28 @@ def test_sanitize_strips_effort_for_haiku_45(): data = {"output_config": {"effort": "high"}} sanitize_vertex_anthropic_output_params(data, "vertex_ai/claude-opus-4-6") assert data["output_config"] == {"effort": "high"} + + +def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_model_cost_map): + """Regression: Fable 5.1 rejects forced tool use, so the vertex map entry + advertises native structured output and ``response_format`` must map to + ``output_format`` instead of the tool-based path's forced tool_choice.""" + config = VertexAIAnthropicConfig() + response_format = { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + } + + result_params = config.map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model="claude-fable-5-1", + drop_params=False, + ) + + assert "output_format" in result_params + assert "tool_choice" not in result_params + assert "tools" not in result_params diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py index b477bf3f406..2ccba2b2055 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -2,6 +2,30 @@ import os import pytest +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, +) + + +@pytest.fixture(autouse=True) +def _hermetic_mcp_server_registry(): + """Restore the singleton ``global_mcp_server_manager``'s registry state around every + test, so entries seeded by one test never leak into another on a shared shard.""" + saved_registry = dict(global_mcp_server_manager.registry) + saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) + saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) + saved_oauth_slots = global_mcp_server_manager._oauth_discovery_slots + try: + yield + finally: + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry.update(saved_registry) + global_mcp_server_manager.config_mcp_servers.clear() + global_mcp_server_manager.config_mcp_servers.update(saved_config_servers) + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.update(saved_tool_mapping) + global_mcp_server_manager._oauth_discovery_slots = saved_oauth_slots + @pytest.fixture(autouse=True) def _hermetic_server_root_path(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py index ac7082c2668..1c65adac4c6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_identity_env.py @@ -23,6 +23,12 @@ MGMT_MODULE = "litellm.proxy.management_endpoints.mcp_management_endpoints" @contextlib.contextmanager def _env_and_reload(**env): saved = {key: os.environ.get(key) for key in env} + utils_module = importlib.import_module(UTILS_MODULE) + mgmt_module = importlib.import_module(MGMT_MODULE) + # Restore pre-reload module attributes afterwards instead of reloading again: + # a reload re-creates the module's classes, breaking exception identity for + # modules that imported them earlier + snapshots = {module: dict(vars(module)) for module in (utils_module, mgmt_module)} def _apply_env(values): for key, value in values.items(): @@ -32,8 +38,8 @@ def _env_and_reload(**env): os.environ[key] = value def _reload(): - utils = importlib.reload(importlib.import_module(UTILS_MODULE)) - mgmt = importlib.reload(importlib.import_module(MGMT_MODULE)) + utils = importlib.reload(utils_module) + mgmt = importlib.reload(mgmt_module) return utils, mgmt try: @@ -41,7 +47,10 @@ def _env_and_reload(**env): yield _reload() finally: _apply_env(saved) - _reload() + for module, snapshot in snapshots.items(): + for key in [key for key in vars(module) if key not in snapshot]: + delattr(module, key) + vars(module).update(snapshot) def test_defaults_used_when_env_unset(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..0480bbc40a7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -214,6 +214,46 @@ class TestExecuteWithMcpClient: assert server.scopes == ["read", "write"] assert server.has_client_credentials is True + async def test_preview_forwards_per_server_timeout_to_client_factory(self, monkeypatch): + """The request's per-server timeout must reach the temporary MCPServer model: + the client factory reads ``server.timeout`` for both the per-request timeout + and the preview's whole-walk listing deadline.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["server"] = kwargs.get("server") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="slow-catalog-server", + url="https://example.com", + timeout=120.5, + ) + + result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation) + + assert result["status"] == "ok" + assert captured["server"].timeout == 120.5 + @pytest.mark.asyncio async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch): """For M2M OAuth servers the incoming Authorization header (which carries @@ -524,6 +564,131 @@ class TestTestToolsList: assert captured["oauth2_headers"] is None assert oauth_call_counter["count"] == 0 + async def test_preview_tools_list_times_out_on_slow_pagination(self, monkeypatch): + """A preview whose upstream paginates past the listing deadline returns a + timeout error instead of holding the request open.""" + monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False) + monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False) + + class SlowClient: + async def list_tools(self, raise_on_error=False): + await asyncio.sleep(1) + return [] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(SlowClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["status"] == "error" + assert result["error"] is True + assert "Timed out listing tools" in result["message"] + + async def test_preview_tools_list_succeeds_within_deadline(self, monkeypatch): + """The preview timeout scope passes a fast listing through untouched.""" + from mcp.types import Tool as MCPTool + + class QuickClient: + async def list_tools(self, raise_on_error=False): + return [MCPTool(name="quick_tool", description="q", inputSchema={})] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(QuickClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["error"] is None + assert result["message"] == "Successfully retrieved tools" + assert [tool["name"] for tool in result["tools"]] == ["quick_tool"] + + async def test_preview_tools_list_honors_per_server_timeout(self, monkeypatch): + """A per-server timeout above the global default extends the preview deadline.""" + monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False) + monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False) + + from mcp.types import Tool as MCPTool + + class SlowConfiguredClient: + timeout = 1.0 + + async def list_tools(self, raise_on_error=False): + await asyncio.sleep(0.2) + return [MCPTool(name="slow_tool", description="s", inputSchema={})] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(SlowConfiguredClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["error"] is None + assert [tool["name"] for tool in result["tools"]] == ["slow_tool"] + async def test_extracts_oauth2_headers(self, monkeypatch): """Ensure oauth2 auth type pulls oauth headers and omits MCP auth header.""" @@ -786,9 +951,7 @@ class TestListToolsRestAPI: they do for a gateway session, never to the bare session key.""" from litellm.constants import UI_SESSION_TOKEN_TEAM_ID - session_auth = UserAPIKeyAuth( - team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" - ) + session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user") admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org") async def fake_reload(user_id): @@ -868,9 +1031,7 @@ class TestListToolsRestAPI: from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import LiteLLM_ObjectPermissionTable - session_auth = UserAPIKeyAuth( - team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" - ) + session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user") scoped_auth = UserAPIKeyAuth( object_permission=LiteLLM_ObjectPermissionTable( object_permission_id="toolset-scope", @@ -952,6 +1113,123 @@ class TestListToolsRestAPI: assert scope_inputs == [session_auth] assert reload_calls == [] + async def test_single_server_response_includes_paginated_upstream_tools( + self, + monkeypatch, + ): + """The REST tools/list path should include tools beyond the upstream first page.""" + import litellm.experimental_mcp_client.client as mcp_client_module + from mcp.types import ListToolsResult, PaginatedRequestParams + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + stub_server = MCPServer( + server_id="server-1", + name="stub", + server_name="stub", + alias="stub", + url="https://example.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "stub"}, + ) + stub_server.available_on_public_internet = True + + mock_transport_ctx = AsyncMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) + mock_transport_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "streamable_http_client", + MagicMock(return_value=mock_transport_ctx), + raising=False, + ) + + mock_session_ctx = AsyncMock() + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=None) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[ + MCPTool( + name="first_page_tool", + description="First page tool", + inputSchema={}, + ) + ], + nextCursor="page-2", + ), + ListToolsResult( + tools=[ + MCPTool( + name="second_page_tool", + description="Second page tool", + inputSchema={}, + ) + ] + ), + ] + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + mock_session_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "ClientSession", + MagicMock(return_value=mock_session_ctx), + raising=False, + ) + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "filter_server_ids_by_ip_with_info", + lambda server_ids, client_ip: (server_ids, 0), + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert set(result.keys()) == {"tools", "error", "message"} + assert [tool.name for tool in result["tools"]] == [ + "first_page_tool", + "second_page_tool", + ] + assert result["error"] is None + assert result["message"] == "Successfully retrieved tools" + + assert mock_session_instance.list_tools.call_count == 2 + second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; a non-admin passing it stays filtered so the REST endpoint can't be used @@ -1153,7 +1431,11 @@ class TestListToolsRestAPI: async def test_aggregate_list_absorbs_one_server_auth_failure(self, monkeypatch): """The multi-server aggregate listing degrades a server whose upstream rejects auth to an empty contribution and still returns the healthy - server's tools with a 200, rather than surfacing a 401.""" + server's tools with a 200, rather than surfacing a 401. The absorbed + server must still show up as a classified per-server outcome so a REST + caller can tell "needs upstream auth" apart from "has no tools".""" + from pydantic import TypeAdapter + from litellm.proxy._experimental.mcp_server.exceptions import ( MCPUpstreamAuthError, ) @@ -1219,6 +1501,11 @@ class TestListToolsRestAPI: assert result["tools"] == ["good-tool"] assert result["error"] is None + wire_body = json.loads(TypeAdapter(dict).dump_json(result)) + assert wire_body["server_outcomes"] == { + "good": {"status": "ok", "tool_count": 1}, + "bad": {"status": "auth_required", "http_status": 401}, + } async def test_name_resolution_finds_server_by_uuid(self, monkeypatch): """When server_id is a name string, it should be resolved to its UUID @@ -3021,9 +3308,7 @@ class TestRestListToolsetFiltering: mock_manager = MagicMock() mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) - mock_manager.resolve_toolset_tool_permissions = AsyncMock( - return_value={"server-a": ["lookup_status"]} - ) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value={"server-a": ["lookup_status"]}) monkeypatch.setattr( rest_endpoints.global_mcp_server_manager, diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py new file mode 100644 index 00000000000..b7bc670c7f8 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -0,0 +1,288 @@ +""" +Tests for restamping the public model on Anthropic Messages streaming chunks. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, + restamp_anthropic_stream_chunk_model, +) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def _message_start_frame(model: str, line_end: str = "\n") -> bytes: + payload = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, + } + return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() + + +def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: + async def _iterator_hook(**_kwargs): + for frame in frames: + yield frame + + proxy_logging_obj = MagicMock() + proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + return proxy_logging_obj + + +def _model_from_frame(frame: bytes | str) -> str: + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + data_line = next(line for line in text.split("\n") if line.startswith("data:")) + return json.loads(data_line[len("data:") :])["message"]["model"] + + +def test_restamps_sse_bytes_frame(): + restamped = restamp_anthropic_stream_chunk_model( + _message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1" + ) + + assert isinstance(restamped, bytes) + assert _model_from_frame(restamped) == "claude-auto-1" + assert b"event: message_start" in restamped + + +def test_restamps_event_dict(): + chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} + + restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2") + + assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}} + assert chunk["message"]["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize( + "chunk", + [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + {"type": "content_block_delta", "delta": {"text": "hi"}}, + {"type": "message_start", "message": "not-a-dict"}, + b"event: message_start\ndata: not-json\n\n", + b"data: [DONE]\n\n", + ], +) +def test_leaves_chunks_without_a_model_untouched(chunk): + assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk + + +@pytest.mark.asyncio +async def test_sse_generator_publishes_requested_model_on_message_start(): + """The message_start event reports the requested model, not the provider's.""" + delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta_frame + + +@pytest.mark.asyncio +async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" + + +def test_restamps_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_emits_coalesced_frames_with_only_message_start_rewritten(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + combined = _message_start_frame("claude-haiku-4-5-20251001") + delta + + emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined) + + assert isinstance(restamper_output, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(delta) + + +def test_ping_frames_keep_the_restamper_armed(): + ping = b'event: ping\ndata: {"type": "ping"}\n\n' + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(ping) == ping + reassembled = restamper.process(frame[:10]) + reassembled += restamper.process(frame[10:]) + + assert _model_from_frame(reassembled) == "claude-auto-1" + + +def test_first_non_ping_event_disarms_the_restamper(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + late_message_start = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(delta) == delta + assert restamper.process(late_message_start) == late_message_start + + +def test_oversized_unterminated_chunk_flushes_unmodified(): + blob = b"data: " + b"x" * 70000 + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(blob) == blob + frame = _message_start_frame("claude-haiku-4-5-20251001") + assert restamper.process(frame) == frame + + +def test_dict_message_start_disarms_after_restamp(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}) + second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}} + + assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}} + assert restamper.process(second) == second + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_message_start_split_across_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert _model_from_frame(joined) == "claude-auto-1" + + +def test_restamps_crlf_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(b"\r\n\r\n") + assert restamper.process(delta) == delta + + +def test_restamps_cr_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert b'"model":"claude-auto-1"' in emitted + assert emitted.endswith(b"\r\r") + + +def test_restamps_crlf_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_flush_returns_restamped_held_tail(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(unterminated) == b"" + flushed = restamper.flush() + + assert b'"model":"claude-auto-1"' in flushed + assert restamper.flush() == b"" + + +def test_flush_disarms_the_restamper(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + frame = _message_start_frame("claude-haiku-4-5-20251001") + + assert restamper.flush() == b"" + assert restamper.process(frame) == frame + + +@pytest.mark.asyncio +async def test_sse_generator_flushes_held_tail_at_end_of_stream(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert b'"model":"claude-auto-1"' in joined + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_crlf_stream(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index 898f9ab1ed7..e5f2a9d95bd 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -26,6 +26,64 @@ def _owners(*backup_paths): CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" +WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" + +CMD_METACHARACTERS = frozenset("&|<>^()") +CMD_PERCENT_GUARD = "%%cd:~,%" + + +def _through_cmd_exe(command): + """The line cmd.exe hands to CreateProcess after reading the apiKeyHelper. + + A `"` toggles cmd's quote state and the metacharacters only act outside it. cmd expands + `%VAR%` even inside quotes, so every `%` has to arrive as the `%%cd:~,%` guard: the first + `%` has no variable name and stays literal, and `%cd:~,%` is a zero length substring of `cd`. + """ + assert not any(CMD_METACHARACTERS & set(run) for run in command.split('"')[::2]), command + assert command.count("%") == 3 * command.count(CMD_PERCENT_GUARD), command + return command.replace(CMD_PERCENT_GUARD, "%") + + +def _through_c_runtime(command_line): + """argv as the Microsoft C runtime builds it for the `lite` executable. + + Outside quotes whitespace ends an argument. A `"` toggles quoting, and inside quotes `""` + is a literal quote. Backslashes are literal unless they run up to a `"`, where each pair + is one backslash and an odd one left over makes the quote literal. + """ + argv = [] + current = None + quoted = False + i = 0 + while i < len(command_line): + ch = command_line[i] + if ch in " \t" and not quoted: + if current is not None: + argv.append(current) + current = None + i += 1 + continue + if current is None: + current = "" + if ch == "\\": + run = len(command_line[i:]) - len(command_line[i:].lstrip("\\")) + before_quote = command_line[i + run : i + run + 1] == '"' + current += "\\" * (run // 2 if before_quote else run) + if before_quote and run % 2: + current += '"' + i += 1 + i += run + elif ch == '"': + if quoted and command_line[i + 1 : i + 2] == '"': + current += '"' + i += 1 + else: + quoted = not quoted + i += 1 + else: + current += ch + i += 1 + return argv if current is None else [*argv, current] @pytest.fixture @@ -199,6 +257,36 @@ class TestApiKeyHelperIsActuallyInvocable: assert "Not authenticated for this server" in result.output + def _windows_argv(self, lite_exe, base_url): + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=lite_exe): + helper = resolve_api_key_helper(base_url, platform="win32") + return _through_c_runtime(_through_cmd_exe(helper)) + + @pytest.mark.parametrize( + ("lite_exe", "base_url"), + [ + (WINDOWS_LITE_EXE, "http://localhost:4000"), + ("C:\\Program Files\\LiteLLM\\lite.EXE", "https://gateway.example.com/?a=1&b=2"), + ("C:\\Users\\u\\Scripts\\lite.EXE", "https://gateway.example.com/team%20a/%7Eproxy"), + ('C:\\odd "dir"\\lite.EXE', "http://localhost:4000/x\\"), + ], + ) + def test_the_windows_command_survives_cmd_exe_and_the_c_runtime(self, lite_exe, base_url): + assert self._windows_argv(lite_exe, base_url) == [lite_exe, "--base-url", base_url, "auth", "print-token"] + + def test_the_windows_command_carries_the_base_url_through_cmd_quoting(self): + stale = CliTokenRecord( + base_url="http://other-proxy.example.com", + key="sk-stale", + timestamp=time.time(), + ) + argv = self._windows_argv(WINDOWS_LITE_EXE, "http://localhost:4000") + with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): + result = CliRunner().invoke(cli, argv[1:]) + + assert argv[0] == WINDOWS_LITE_EXE + assert "Not authenticated for this server" in result.output + class TestConflictingOwnersOfTheSettingsFile: """Both `lite up` and `lite autoroute up` restore a backup when they stop. diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 053c90c36b4..c78bdfa75b1 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -225,6 +225,32 @@ class TestResolveApiKeyHelper: with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): resolve_api_key_helper("http://localhost:4000") + def test_windows_quotes_for_cmd_exe_instead_of_posix_sh(self, monkeypatch): + """cmd.exe takes a single quote literally, so a POSIX-quoted backslashed path is unrunnable.""" + lite_exe = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" + monkeypatch.setattr(shutil, "which", lambda name: lite_exe) + + helper = resolve_api_key_helper("https://gateway.example.com", platform="win32") + + assert helper == f'"{lite_exe}" "--base-url" "https://gateway.example.com" "auth" "print-token"' + + def test_windows_keeps_a_spaced_path_and_a_metacharacter_url_as_single_tokens(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "C:\\Program Files\\LiteLLM\\lite.EXE") + + helper = resolve_api_key_helper("https://gateway.example.com/?a=1&b=2", platform="win32") + + assert helper == ( + '"C:\\Program Files\\LiteLLM\\lite.EXE" "--base-url" "https://gateway.example.com/?a=1&b=2" ' + '"auth" "print-token"' + ) + + def test_non_windows_platforms_keep_posix_quoting(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + + helper = resolve_api_key_helper("http://example.com/path; rm -rf /", platform="darwin") + + assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" + def _make_ctx(base_url): return click.Context(click.Command("test"), obj={"base_url": base_url}) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_alice.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_alice.py new file mode 100644 index 00000000000..fd2e86ccde8 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_alice.py @@ -0,0 +1,614 @@ +import json +import os +from copy import deepcopy +from unittest.mock import AsyncMock + +import httpx +import pytest +from httpx import Request, Response + +import litellm +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy.guardrails.guardrail_hooks.alice.alice import ( + GUARDRAIL_NAME, + AliceGuardrail, + AliceGuardrailMissingSecrets, + _json_safe, +) +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 + + +def _guardrail(**overrides: object) -> AliceGuardrail: + params: dict[str, object] = {"api_key": "test-key", "guardrail_name": "alice", "event_hook": "pre_call"} + params.update(overrides) + return AliceGuardrail(**params) + + +def _verdict(payload: dict[str, object], status_code: int = 200) -> Response: + return Response( + status_code=status_code, + json=payload, + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + + +def test_alice_guardrail_config(monkeypatch: pytest.MonkeyPatch): + """Should register through init_guardrails_v2 like any other provider.""" + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setenv("ALICE_API_KEY", "test-key") + + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "alice", + "litellm_params": {"guardrail": "alice", "mode": "pre_call", "default_on": True}, + } + ], + config_file_path="", + ) + + registered = [cb for cb in litellm.callbacks if isinstance(cb, AliceGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "alice" + + +class TestAliceGuardrailInitialization: + def setup_method(self): + for key in ("ALICE_API_KEY", "ALICE_API_BASE"): + os.environ.pop(key, None) + + def test_missing_api_key_raises(self): + with pytest.raises(AliceGuardrailMissingSecrets, match="API key"): + AliceGuardrail(guardrail_name="alice", event_hook="pre_call") + + def test_reads_credentials_from_environment(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ALICE_API_KEY", "env-key") + monkeypatch.setenv("ALICE_API_BASE", "https://env.alice.test") + + guardrail = AliceGuardrail(guardrail_name="alice", event_hook="pre_call") + + assert guardrail.alice_api_key == "env-key" + assert guardrail.api_base == "https://env.alice.test/v2/evaluate/litellm" + + def test_defaults_the_api_base(self): + assert _guardrail().api_base == "https://api.alice.io/v2/evaluate/litellm" + + def test_trailing_slash_does_not_double_up(self): + assert _guardrail(api_base="https://api.alice.io/").api_base == ("https://api.alice.io/v2/evaluate/litellm") + + +class TestAliceForwarding: + """The hook's arguments cross the wire as they were received — nothing selected, nothing + renamed — except the caller's raw credentials, which are stripped before request_data is + serialized (see TestAliceCredentialStripping).""" + + @pytest.mark.asyncio + async def test_forwards_the_hook_arguments_verbatim(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + inputs = {"texts": ["hello"], "structured_messages": [{"role": "user", "content": "hello"}]} + request_data = {"model": "gpt-4o", "metadata": {"user_api_key_alias": "payments-bot"}} + # Snapshot before the call: @log_guardrail_information writes its own entry into + # request_data["metadata"] afterwards, so the original is no longer what was sent. + sent_inputs = deepcopy(inputs) + sent_request_data = deepcopy(request_data) + + await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request") + + body = guardrail.async_handler.post.call_args.kwargs["json"] + assert body["input_type"] == "request" + assert body["inputs"] == sent_inputs + assert body["request_data"] == sent_request_data + + @pytest.mark.asyncio + async def test_sends_the_credential(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request") + + assert guardrail.async_handler.post.call_args.kwargs["headers"]["af-api-key"] == "test-key" + + @pytest.mark.asyncio + async def test_marks_a_completion_as_a_response(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + + await guardrail.apply_guardrail(inputs={"texts": ["answer"]}, request_data={}, input_type="response") + + assert guardrail.async_handler.post.call_args.kwargs["json"]["input_type"] == "response" + + @pytest.mark.asyncio + async def test_nothing_selectable_reaches_no_evaluation(self): + """No texts, images, tools, tool_calls, or structured_messages: genuinely nothing to send.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock() + + result = await guardrail.apply_guardrail(inputs={"texts": []}, request_data={}, input_type="request") + + assert result == {"texts": []} + guardrail.async_handler.post.assert_not_called() + + @pytest.mark.asyncio + async def test_tool_calls_only_still_reaches_alice(self): + """A batch with empty texts but populated tool_calls is still a selection decision Alice + should make, not the plugin — see the class docstring.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + inputs = {"texts": [], "tool_calls": [{"id": "call_1", "function": {"name": "get_weather"}}]} + + await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + guardrail.async_handler.post.assert_called_once() + assert guardrail.async_handler.post.call_args.kwargs["json"]["inputs"]["tool_calls"] == inputs["tool_calls"] + + @pytest.mark.asyncio + async def test_images_only_still_reaches_alice(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + + await guardrail.apply_guardrail( + inputs={"texts": [], "images": ["data:image/png;base64,abc"]}, request_data={}, input_type="request" + ) + + guardrail.async_handler.post.assert_called_once() + + @pytest.mark.asyncio + async def test_structured_messages_only_still_reaches_alice(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + + await guardrail.apply_guardrail( + inputs={"texts": [], "structured_messages": [{"role": "user", "content": []}]}, + request_data={}, + input_type="request", + ) + + guardrail.async_handler.post.assert_called_once() + + @pytest.mark.asyncio + async def test_makes_exactly_one_attempt(self): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock(side_effect=httpx.ConnectError("refused")) + + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request") + + assert guardrail.async_handler.post.call_count == 1 + + +class TestAliceCredentialStripping: + """request_data's raw-credential keys never leave the process.""" + + @pytest.mark.asyncio + async def test_secret_fields_and_api_key_are_stripped(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + request_data = { + "model": "gpt-4o", + "api_key": "sk-forwarded-provider-secret", + "secret_fields": {"raw_headers": {"authorization": "Bearer caller-virtual-key"}}, + "metadata": {"user_api_key_alias": "payments-bot"}, + } + + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request") + + sent_request_data = guardrail.async_handler.post.call_args.kwargs["json"]["request_data"] + assert "secret_fields" not in sent_request_data + assert "api_key" not in sent_request_data + assert sent_request_data == {"model": "gpt-4o", "metadata": {"user_api_key_alias": "payments-bot"}} + + @pytest.mark.asyncio + async def test_nested_credentials_are_stripped_at_every_depth(self): + """Shaped after a real captured Claude Code payload: the caller's Authorization/x-api-key + lives under several independent nesting paths, none of which are the root.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + request_data = { + "model": "claude-3-5-sonnet", + "secret_fields": {"raw_headers": {"authorization": "Bearer caller-virtual-key"}}, + "provider_specific_header": {"extra_headers": {"authorization": "sk-ant-oat01-nested-oauth"}}, + "proxy_server_request": { + "url": "/v1/messages", + "headers": {"authorization": "Bearer inbound-caller-secret", "x-request-id": "req-1"}, + "body": { + "model": "claude-3-5-sonnet", + "metadata": {"headers": {"authorization": "Bearer body-metadata-secret"}}, + }, + }, + "metadata": { + "user_api_key_alias": "payments-bot", + "headers": {"authorization": "Bearer metadata-secret"}, + "requester_metadata": {"headers": {"authorization": "Bearer requester-metadata-secret"}}, + }, + "litellm_metadata": {"headers": {"authorization": "Bearer litellm-metadata-secret"}}, + } + + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request") + + posted_body = guardrail.async_handler.post.call_args.kwargs["json"] + serialized = json.dumps(posted_body) + assert "authorization" not in serialized.lower() + assert "caller-virtual-key" not in serialized + assert "nested-oauth" not in serialized + assert "inbound-caller-secret" not in serialized + assert "body-metadata-secret" not in serialized + assert "metadata-secret" not in serialized + assert "requester-metadata-secret" not in serialized + assert "litellm-metadata-secret" not in serialized + + sent_request_data = posted_body["request_data"] + assert sent_request_data["model"] == "claude-3-5-sonnet" + assert sent_request_data["proxy_server_request"]["url"] == "/v1/messages" + assert "headers" not in sent_request_data["proxy_server_request"] + assert sent_request_data["proxy_server_request"]["body"]["model"] == "claude-3-5-sonnet" + assert "headers" not in sent_request_data["proxy_server_request"]["body"]["metadata"] + assert sent_request_data["metadata"]["user_api_key_alias"] == "payments-bot" + assert "headers" not in sent_request_data["metadata"] + assert "requester_metadata" in sent_request_data["metadata"] + assert "headers" not in sent_request_data["metadata"]["requester_metadata"] + assert "headers" not in sent_request_data["litellm_metadata"] + assert "secret_fields" not in sent_request_data + assert "provider_specific_header" not in sent_request_data + + @pytest.mark.asyncio + async def test_the_original_request_data_is_not_mutated(self): + """Stripping must only affect the outbound copy — api_key still has to reach the + provider, and secret_fields still has to reach the rest of the request pipeline.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + request_data = {"api_key": "sk-forwarded-provider-secret", "secret_fields": {"raw_headers": {}}} + + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data=request_data, input_type="request") + + assert request_data["api_key"] == "sk-forwarded-provider-secret" + assert request_data["secret_fields"] == {"raw_headers": {}} + + +class TestAliceVerdicts: + @pytest.mark.asyncio + async def test_allow_leaves_the_inputs_untouched(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "ALLOW", "categories": []})) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_block_surfaces_the_policy_message(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict( + { + "verdict": "BLOCK", + "categories": ["self_harm"], + "correlation_id": "c1", + "message": "Blocked by your organization's policy", + } + ) + ) + + with pytest.raises(GuardrailRaisedException) as error: + await guardrail.apply_guardrail(inputs={"texts": ["bad"]}, request_data={}, input_type="request") + + assert "Blocked by your organization's policy" in str(error.value) + + @pytest.mark.asyncio + async def test_block_without_a_message_still_blocks(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "BLOCK", "categories": []})) + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail(inputs={"texts": ["bad"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_mask_substitutes_by_position(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict( + { + "verdict": "MASK", + "categories": ["pii"], + "replacements": [{"index": 1, "text": "my ssn is ***"}], + } + ) + ) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["untouched", "my ssn is 123-45-6789"]}, + request_data={}, + input_type="request", + ) + + assert result["texts"] == ["untouched", "my ssn is ***"] + + @pytest.mark.asyncio + async def test_mask_that_lands_nowhere_blocks(self): + """A mask that wrote nothing would let the text through under a verdict that said not to.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict({"verdict": "MASK", "categories": [], "replacements": [{"index": 9, "text": "***"}]}) + ) + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_mask_with_no_replacements_blocks(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(return_value=_verdict({"verdict": "MASK", "categories": []})) + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_mask_with_one_invalid_replacement_blocks_entirely(self): + """A mixed valid/invalid replacement list must not let the valid half through: that + would leave the content named by the invalid entry unmasked while looking like success.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict( + { + "verdict": "MASK", + "categories": ["pii"], + "replacements": [{"index": 0, "text": "***"}, {"index": 9, "text": "***"}], + } + ) + ) + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"]}, request_data={}, input_type="request" + ) + + @pytest.mark.asyncio + async def test_mask_leaves_structured_messages_identical(self): + """A new structured_messages object makes the translation layer skip the texts write-back.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict({"verdict": "MASK", "categories": [], "replacements": [{"index": 0, "text": "***"}]}) + ) + messages = [{"role": "user", "content": "secret"}] + + result = await guardrail.apply_guardrail( + inputs={"texts": ["secret"], "structured_messages": messages}, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] is messages + + @pytest.mark.asyncio + async def test_detect_allows_and_leaves_the_text_alone(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=_verdict({"verdict": "DETECT", "categories": ["profanity"], "correlation_id": "c1"}) + ) + + result = await guardrail.apply_guardrail(inputs={"texts": ["mild"]}, request_data={}, input_type="request") + + assert result["texts"] == ["mild"] + + +class TestAliceUnreachable: + @pytest.mark.parametrize( + "failure", + [ + pytest.param({"side_effect": httpx.ConnectError("refused")}, id="connect-error"), + pytest.param({"return_value": _verdict({"verdict": "MAYBE"})}, id="unrecognized-verdict"), + pytest.param({"return_value": _verdict({})}, id="no-verdict"), + ], + ) + @pytest.mark.asyncio + async def test_fails_closed_by_default(self, failure: dict): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock(**failure) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_fails_open_when_configured(self): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock(side_effect=httpx.ConnectError("refused")) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + +class TestAliceTransportFailures: + """Every path out of the HTTP call, since each decides whether traffic flows unscreened.""" + + @pytest.mark.asyncio + async def test_a_timeout_is_unreachable(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + side_effect=litellm.exceptions.Timeout(message="slow", model="gpt-4o", llm_provider="openai") + ) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.parametrize("status", [500, 502, 503, 504]) + @pytest.mark.asyncio + async def test_upstream_5xx_is_unreachable(self, status: int): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + response=_verdict({}, status_code=status), + ) + ) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_a_500_fails_closed_by_default(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + response=_verdict({}, status_code=500), + ) + ) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_a_4xx_is_not_treated_as_unreachable(self): + """A rejected credential is our misconfiguration, not an outage — it must not fail open.""" + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "unauthorized", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + response=_verdict({}, status_code=401), + ) + ) + + with pytest.raises(httpx.HTTPStatusError): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_a_non_object_body_fails_closed_by_default(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + json=["not", "an", "object"], + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_a_non_object_body_fails_open_when_configured(self): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + json=["not", "an", "object"], + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_malformed_json_fails_closed_by_default(self): + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + content=b"not json", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_malformed_json_fails_open_when_configured(self): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + content=b"not json", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + @pytest.mark.asyncio + async def test_an_undecodable_body_fails_closed_by_default(self): + """UnicodeDecodeError is a sibling of JSONDecodeError under ValueError, not a subclass.""" + guardrail = _guardrail() + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + content=b"\xff\xfe not utf-8", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + with pytest.raises(GuardrailRaisedException, match="unavailable"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + @pytest.mark.asyncio + async def test_an_undecodable_body_fails_open_when_configured(self): + guardrail = _guardrail(unreachable_fallback="fail_open") + guardrail.async_handler.post = AsyncMock( + return_value=Response( + status_code=200, + content=b"\xff\xfe not utf-8", + request=Request("POST", "https://api.alice.io/v2/evaluate/litellm"), + ) + ) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert result["texts"] == ["hello"] + + +class TestAliceSerialization: + """`request_data` carries live objects, so it cannot be posted as it stands.""" + + def test_drops_what_cannot_serialize_and_keeps_the_rest(self): + class Span: + pass + + result = _json_safe({"model": "x", "metadata": {"span": Span(), "user": "u1"}, "n": 1}) + + assert result == {"model": "x", "metadata": {"span": None, "user": "u1"}, "n": 1} + + def test_survives_a_cycle(self): + data: dict = {"a": 1} + data["self"] = data + + assert _json_safe(data) == {"a": 1, "self": None} + + def test_drops_a_model_that_will_not_dump(self): + class Stubborn: + def model_dump(self, mode: str = "python") -> dict: + raise RuntimeError("cannot serialise") + + assert _json_safe({"m": Stubborn()}) == {"m": None} + + def test_drops_a_bare_unserialisable_value(self): + class Span: + pass + + assert _json_safe(Span()) is None + + def test_dumps_pydantic_models(self): + from pydantic import BaseModel + + class Model(BaseModel): + name: str + + assert _json_safe({"m": Model(name="x")}) == {"m": {"name": "x"}} + + +def test_config_model_is_exposed_for_the_ui(): + config_model = AliceGuardrail.get_config_model() + + assert config_model is not None + assert config_model.ui_friendly_name() == "Alice" + + +def test_guardrail_name_constant(): + assert GUARDRAIL_NAME == "alice" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 235a5c0c09b..953e3de1519 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5524,7 +5524,9 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream scan used to raise after SSE headers were flushed, so the client saw a silently truncated stream. The unified hook must emit the chat in-stream - error frame instead.""" + error frame instead. The finish chunk is withheld while the end-of-stream + scan runs, so on a block it is dropped rather than relayed before the + frame.""" from litellm.llms import load_guardrail_translation_mappings from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( unified_guardrail as unified_module, @@ -5582,11 +5584,211 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca finally: unified_module.endpoint_guardrail_translation_mappings = None - assert len(out) == 3 + assert len(out) == 2 assert isinstance(out[0], ModelResponseStream) + assert out[0].choices[0].finish_reason is None frame = out[-1] assert isinstance(frame, bytes) payload = json.loads(frame.decode()[len("data: ") :]) assert payload["error"]["message"] == "Violated guardrail policy" assert payload["error"]["code"] == "400" assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" + + +def _responses_stream_events() -> list: + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + deltas = [ + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_lit6457", + output_index=0, + content_index=0, + delta=part, + ) + for part in ("Hello", " world") + ] + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_lit6457", + created_at=1234567890, + model="gpt-4o", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_lit6457", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello world"}], + } + ], + ), + ) + return [*deltas, completed] + + +@pytest.mark.asyncio +async def test_responses_api_stream_scans_output_and_replays_buffered_events(): + """Streamed /v1/responses events must be scanned via the unified translation + layer, not fed to stream_chunk_builder (which raises APIError on them).""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-responses-stream", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + stream_events = _responses_stream_events() + order = [] + yielded = [] + + async def record_scan(*args, **kwargs): + order.append("scan") + return {"action": "NONE", "assessments": [], "outputs": []} + + async def mock_stream(): + for event in stream_events: + yield event + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=mock_stream(), + request_data={"model": "gpt-4o", "input": "hi"}, + ): + order.append("chunk") + yielded.append(chunk) + + assert order == ["scan", "chunk", "chunk", "chunk"] + assert len(yielded) == len(stream_events) + assert all(emitted is original for emitted, original in zip(yielded, stream_events)) + + +def _responses_failed_stream_events() -> list: + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseFailedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + deltas = [ + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_lit6457_failed", + output_index=0, + content_index=0, + delta=part, + ) + for part in ("Hello", " world") + ] + failed = ResponseFailedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_lit6457_failed", + created_at=1234567890, + model="gpt-4o", + object="response", + status="failed", + output=[], + ), + ) + return [*deltas, failed] + + +@pytest.mark.asyncio +async def test_responses_api_failed_stream_scans_delta_text_before_replay(): + """A responses stream that dies mid-generation carries its text only in delta + events; the end-of-stream scan must still see that text instead of skipping + on an empty assembled string and replaying the buffer unmoderated.""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-responses-failed-stream", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + stream_events = _responses_failed_stream_events() + order = [] + scan_payloads = [] + yielded = [] + + async def record_scan(*args, **kwargs): + order.append("scan") + scan_payloads.append(str(args) + str(kwargs)) + return {"action": "NONE", "assessments": [], "outputs": []} + + async def mock_stream(): + for event in stream_events: + yield event + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=mock_stream(), + request_data={"model": "gpt-4o", "input": "hi"}, + ): + order.append("chunk") + yielded.append(chunk) + + assert order == ["scan", "chunk", "chunk", "chunk"] + assert "Hello world" in scan_payloads[0] + assert len(yielded) == len(stream_events) + assert all(emitted is original for emitted, original in zip(yielded, stream_events)) + + +@pytest.mark.asyncio +async def test_apply_guardrail_debug_log_masks_signed_request_headers(): + import logging + + from litellm._logging import verbose_proxy_logger + + session_token = "FakeSessionTokenValueThatMustNeverAppearInLogs1234567890" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_access_key_id="ASIAFAKEACCESSKEYID1", + aws_secret_access_key="fakeSecretAccessKeyForSigning", + aws_session_token=session_token, + aws_region_name="us-east-1", + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"action": "NONE", "outputs": []} + + captured_records: list[logging.LogRecord] = [] + + class _RecordingHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + captured_records.append(record) + + handler = _RecordingHandler(level=logging.DEBUG) + previous_level = verbose_proxy_logger.level + verbose_proxy_logger.addHandler(handler) + verbose_proxy_logger.setLevel(logging.DEBUG) + try: + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={}, + ) + finally: + verbose_proxy_logger.removeHandler(handler) + verbose_proxy_logger.setLevel(previous_level) + + rendered_messages = [record.getMessage() for record in captured_records] + header_lines = [message for message in rendered_messages if "headers:" in message] + assert header_lines, "expected the signed-request debug line to be logged" + assert any("X-Amz-Security-Token" in message for message in header_lines) + assert all(session_token not in message for message in rendered_messages) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py new file mode 100644 index 00000000000..42bdf41bc88 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py @@ -0,0 +1,327 @@ +""" +Regression tests for blocking an OpenAI-format streaming response from the +unified guardrail post-call streaming iterator hook. + +When a guardrail's ``apply_guardrail`` raises ``ModifyResponseException`` +while (or at the end of) a chat completions or Responses API stream is being +relayed, the hook must emit a well-formed SSE termination sequence carrying +the block message - NOT a bare ``data: {"error": ...}`` blob that surfaces as +an HTTP 500 error frame and truncates the stream. +""" + +import json +from typing import Any, AsyncGenerator, Dict, Literal, Optional, Tuple, Union + +import pytest + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.utils import ( + Delta, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, +) + +BLOCK_MESSAGE = "This response was replaced by policy." + +JsonPayload = Dict[str, object] +StreamChunk = Union[ModelResponseStream, JsonPayload, bytes] + + +class _BlockingGuardrail(CustomGuardrail): + """Mock guardrail that always blocks response scans by raising ModifyResponseException.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-5.4-mini", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + +class _PassingGuardrail(CustomGuardrail): + """Mock guardrail that always lets response scans through unchanged.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return inputs + + +def _chat_chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-live", + created=1724900000, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + +async def _chat_stream(end: bool) -> AsyncGenerator[ModelResponseStream, None]: + yield _chat_chunk(Delta(role="assistant", content="This ")) + for text in ["is ", "the ", "original ", "answer."]: + yield _chat_chunk(Delta(content=text)) + if end: + yield _chat_chunk(Delta(), finish_reason="stop") + + +async def _responses_stream(end: bool) -> AsyncGenerator[JsonPayload, None]: + original_text = "This is the original answer." + response_envelope = {"id": "resp_live", "model": "gpt-5.4-mini", "status": "in_progress", "output": []} + yield {"type": "response.created", "response": response_envelope} + yield {"type": "response.in_progress", "response": response_envelope} + yield { + "type": "response.output_item.added", + "output_index": 0, + "item": {"id": "msg_orig", "type": "message", "role": "assistant", "content": []}, + } + yield { + "type": "response.content_part.added", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + } + for delta in ["This ", "is ", "the ", "original ", "answer."]: + yield { + "type": "response.output_text.delta", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "delta": delta, + } + yield { + "type": "response.output_text.done", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "text": original_text, + } + if end: + yield { + "type": "response.completed", + "response": { + "id": "resp_live", + "model": "gpt-5.4-mini", + "status": "completed", + "output": [ + { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": original_text, "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + }, + } + + +async def _run_hook( + route: str, + stream: AsyncGenerator[Union[ModelResponseStream, JsonPayload], None], + sampling_rate: int = 1, + end_of_stream_only: bool = False, + buffer_until_moderated: bool = False, + blocks: bool = True, +) -> Tuple[StreamChunk, ...]: + guardrail = ( + _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") + if blocks + else _PassingGuardrail(guardrail_name="test-passing-guardrail", event_hook="post_call") + ) + guardrail.streaming_sampling_rate = sampling_rate + guardrail.streaming_end_of_stream_only = end_of_stream_only + guardrail.streaming_buffer_until_moderated = buffer_until_moderated + + unified_guardrail = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route=route) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + + return tuple( + [ + chunk + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=stream, + request_data=request_data, + ) + ] + ) + + +def _sse_payloads(collected: Tuple[StreamChunk, ...]) -> Tuple[JsonPayload, ...]: + return tuple( + json.loads(line[len("data:") :].strip()) + for chunk in collected + if isinstance(chunk, bytes) + for block in chunk.decode().split("\n\n") + for line in block.strip().split("\n") + if line.startswith("data:") + ) + + +def _assert_no_error_frame(collected: Tuple[StreamChunk, ...]) -> None: + raw = "".join(chunk.decode() for chunk in collected if isinstance(chunk, bytes)) + assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}" + + +@pytest.mark.asyncio +async def test_chat_pre_stream_block_emits_standalone_completion(): + """Block on the first chunk: a standalone completion opens with a role delta + and ends with finish_reason content_filter.""" + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False)) + _assert_no_error_frame(collected) + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant", "content": BLOCK_MESSAGE} + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_chat_mid_stream_block_continues_the_completion(): + """Regression for the LIT-6496 500 error frame: after chunks were already + forwarded, the block continues the same completion id and terminates with + finish_reason content_filter instead of raising into an error blob.""" + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False), sampling_rate=5) + _assert_no_error_frame(collected) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + assert forwarded, "original chunks should have streamed before the block" + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + assert all(payload["id"] == "chatcmpl-live" for payload in payloads), ( + "block chunks must continue the in-progress completion, not start a new one" + ) + assert payloads[0]["choices"][0]["delta"] == {"content": BLOCK_MESSAGE} + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_chat_end_of_stream_block_terminates_cleanly(): + """Regression for bugbot's finish-ordering finding: in end_of_stream_only + mode the original finish chunk must be withheld until moderation decides, + so a block's content_filter finish is the only stream terminator a client + ever sees - never policy text trailing after finish_reason stop.""" + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True) + _assert_no_error_frame(collected) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + assert forwarded, "content chunks still stream to the client before end-of-stream moderation" + assert all(choice.finish_reason is None for chunk in forwarded for choice in chunk.choices), ( + "the original finish chunk must be withheld until moderation decides" + ) + payloads = _sse_payloads(collected) + assert BLOCK_MESSAGE in json.dumps(payloads) + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_chat_end_of_stream_pass_releases_withheld_finish_chunk(): + """When end-of-stream moderation passes, the withheld finish chunk is + released so a clean stream still terminates normally.""" + collected = await _run_hook( + "/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True, blocks=False + ) + assert not [chunk for chunk in collected if isinstance(chunk, bytes)], ( + "a clean stream must carry no synthetic block frames" + ) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + finish_reasons = [choice.finish_reason for chunk in forwarded for choice in chunk.choices] + assert finish_reasons[-1] == "stop", "the withheld finish chunk must be released after moderation passes" + assert all(reason is None for reason in finish_reasons[:-1]) + + +@pytest.mark.asyncio +async def test_responses_buffered_block_emits_full_event_sequence(): + """Buffered moderation blocks before anything streams: a complete synthetic + Responses stream from response.created through response.completed carrying + the block message, with the original content never released.""" + collected = await _run_hook("/v1/responses", _responses_stream(end=True), buffer_until_moderated=True) + _assert_no_error_frame(collected) + assert not [chunk for chunk in collected if isinstance(chunk, dict)], ( + "buffered original chunks must never be released after a block" + ) + payloads = _sse_payloads(collected) + event_types = [payload["type"] for payload in payloads] + assert event_types[0] == "response.created" + assert "response.output_text.delta" in event_types + assert event_types[-1] == "response.completed" + completed = payloads[-1]["response"] + assert completed["status"] == "completed" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + assert "original answer" not in json.dumps(payloads) + + +@pytest.mark.asyncio +async def test_responses_mid_stream_block_continues_the_response(): + """Regression for the LIT-6496 500 error frame and bugbot's unclosed-item + finding: after events were already forwarded, the block first closes the + output item still open on the wire, then appends the replacement item under + the same response id, and closes with response.completed - never a second + response.created and never a completed response with an item left open.""" + collected = await _run_hook("/v1/responses", _responses_stream(end=False)) + _assert_no_error_frame(collected) + forwarded = [chunk for chunk in collected if isinstance(chunk, dict)] + forwarded_types = [chunk["type"] for chunk in forwarded] + assert "response.created" in forwarded_types, "original events should have streamed before the block" + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + block_types = [payload["type"] for payload in payloads] + assert "response.created" not in block_types, "a mid-stream block must not restart the response" + assert block_types[-1] == "response.completed" + + all_events = forwarded + list(payloads) + opened = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.added") + closed = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.done") + assert opened == closed, "every output item opened on the stream must be closed before response.completed" + original_done_position = block_types.index("response.output_item.done") + block_item_position = block_types.index("response.output_item.added") + assert original_done_position < block_item_position, ( + "the in-progress original item must be closed before the block item is appended" + ) + assert payloads[original_done_position]["item"]["id"] == "msg_orig" + assert payloads[block_item_position]["output_index"] == 1, ( + "the block item must continue after the original output item" + ) + completed = payloads[-1]["response"] + assert completed["id"] == "resp_live" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + + +@pytest.mark.asyncio +async def test_responses_end_of_stream_block_reports_original_usage(): + collected = await _run_hook("/v1/responses", _responses_stream(end=True), end_of_stream_only=True) + _assert_no_error_frame(collected) + forwarded_types = [chunk["type"] for chunk in collected if isinstance(chunk, dict)] + assert "response.completed" not in forwarded_types, ( + "the original terminal event must be withheld and replaced by the block sequence" + ) + payloads = _sse_payloads(collected) + completed = payloads[-1]["response"] + assert payloads[-1]["type"] == "response.completed" + assert completed["id"] == "resp_live" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 0b32558a00a..8cad1c634a9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1844,7 +1844,8 @@ class TestStreamingHttpErrorFrames: out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert out[:2] == chunks + assert out[0] == chunks[0] + assert chunks[1] not in out frame = out[-1] assert isinstance(frame, bytes) text = frame.decode() diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9b2117b7647..9511732fd50 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -917,7 +917,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): "Content-Type": "application/json", "Authorization": "Bearer test-api-key-789", } - mock_request_instance.prepare.return_value = Mock() + mock_request_instance.prepare.return_value = Mock( + headers=mock_request_instance.headers + ) mock_aws_request.return_value = mock_request_instance await guardrail_hook.make_bedrock_api_request( diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 957f9fde645..1697b77b99a 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,7 +1,8 @@ import logging import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from itertools import chain +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, call @@ -38,6 +39,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( get_groups, get_users, get_service_provider_config, + merge_placeholder, patch_group, patch_team_membership, patch_user, @@ -52,6 +54,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIMMember, SCIMPatchOp, SCIMPatchOperation, + SCIMPlaceholderMergeResult, SCIMServiceProviderConfig, SCIMUser, SCIMUserEmail, @@ -778,13 +781,17 @@ async def test_handle_existing_user_by_email_without_teams_preserves_memberships "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=None), ) - mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper - "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", - AsyncMock(), + mock_team_member_add = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) ) - mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper - "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", - AsyncMock(), + mock_team_member_delete = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) ) new_user_request = NewUserRequest( @@ -1645,6 +1652,25 @@ async def test_update_group_e2e(mocker): ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team) +def _rows_by_exact_id( + user_row: Callable[[Mapping[str, str]], LiteLLM_UserTable | MagicMock | None], +) -> Callable[..., tuple[LiteLLM_UserTable | MagicMock, ...]]: + """``find_many`` stand-in for the classifier's cross-field read on a table where a + member value only ever matches as an exact ``user_id``.""" + + def rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable | MagicMock, ...]: + clauses: Final = where["OR"] + assert isinstance(clauses, list) + found: Final = tuple(user_row(clause) for clause in clauses if "user_id" in clause) + return tuple(row for row in found if row is not None) + + return rows + + +def _user_row_for(where: Mapping[str, str]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=where["user_id"]) + + @pytest.mark.asyncio async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): """ @@ -1696,9 +1722,8 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -1782,9 +1807,8 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-3 and new-user-4 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -1853,9 +1877,8 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1943,9 +1966,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -2013,9 +2035,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -3121,8 +3142,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # new-user already exists in the DB - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="new-user")) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(mocker.MagicMock(user_id="new-user"),)) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3415,8 +3435,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): ) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for)) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3509,8 +3528,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock ) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for)) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3640,8 +3658,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker): prisma_client = mocker.MagicMock() prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-3")) - prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(LiteLLM_UserTable(user_id="user-3"),)) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3733,12 +3750,14 @@ def _member_resolution_prisma( starts folding it, fails here instead of passing. A caller that must know which accounts match rather than merely how many - passes take=None, so an unbounded read returns every match. + passes take=None, so an unbounded read returns every match. The row keyed by + the value comes last, the order a bounded read is least prepared for, since + the database promises no order at all. """ clauses: Final = where["OR"] assert isinstance(clauses, list) fields: Final = tuple(next(iter(clause)) for clause in clauses) - assert fields == ("sso_user_id", "user_email"), fields + assert fields in (("user_id", "sso_user_id", "user_email"), ("sso_user_id", "user_email")), fields def comparison(clause: Mapping[str, object]) -> tuple[str, bool]: """The needle and whether production asked for a case-insensitive compare, @@ -3749,8 +3768,9 @@ def _member_resolution_prisma( assert isinstance(criterion, dict), criterion return criterion["equals"], criterion.get("mode") == "insensitive" - sso_needle, sso_insensitive = comparison(clauses[0]) - email_needle, email_insensitive = comparison(clauses[1]) + by_field: Final = dict(zip(fields, (comparison(clause) for clause in clauses))) + sso_needle, sso_insensitive = by_field["sso_user_id"] + email_needle, email_insensitive = by_field["user_email"] def same(stored: str, needle: str, insensitive: bool) -> bool: return stored.casefold() == needle.casefold() if insensitive else stored == needle @@ -3768,6 +3788,11 @@ def _member_resolution_prisma( if same(email, email_needle, email_insensitive) for user_id in user_ids ), + ( + user_id + for user_id in users + if "user_id" in by_field and same(user_id, by_field["user_id"][0], by_field["user_id"][1]) + ), ) ) found: Final = tuple(dict.fromkeys(matched)) @@ -4452,9 +4477,11 @@ async def test_create_group_applies_default_team_params( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", AsyncMock(return_value=_member_resolution_prisma(mocker, users=set(), teams=set())), ) - new_team_mock = mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group - "litellm.proxy.management_endpoints.scim.scim_v2.new_team", - AsyncMock(return_value=mocker.MagicMock()), + new_team_mock = ( + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group + "litellm.proxy.management_endpoints.scim.scim_v2.new_team", + AsyncMock(return_value=mocker.MagicMock()), + ) ) mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_team_to_scim_group", @@ -4611,9 +4638,15 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups def _identity_lookup(value: str) -> object: - """The single cross-field lookup the classifier is expected to issue.""" + """The single cross-field lookup the classifier is expected to issue per member.""" return call( - where={"OR": [{"sso_user_id": value}, {"user_email": {"equals": value, "mode": "insensitive"}}]}, + where={ + "OR": [ + {"user_id": value}, + {"sso_user_id": value}, + {"user_email": {"equals": value, "mode": "insensitive"}}, + ] + }, take=2, ) @@ -4903,9 +4936,7 @@ async def test_process_group_patch_remove_by_the_id_the_directory_added_with( @pytest.mark.asyncio -async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id( - mocker, scim_upsert_user_enabled -): +async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id(mocker, scim_upsert_user_enabled): """An earlier release put unmatched ids on the roster verbatim, so a remove has to keep clearing the id as written even once it also resolves.""" patch_ops = SCIMPatchOp( @@ -4916,7 +4947,10 @@ async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_liter team_id="parent-group", team_alias="Parent Group", members=[], - members_with_roles=[Member(user_id="legacy@example.com", role="user"), Member(user_id="keep-user", role="user")], + members_with_roles=[ + Member(user_id="legacy@example.com", role="user"), + Member(user_id="keep-user", role="user"), + ], ) _, final_members, _ = await _process_group_patch_operations( @@ -5081,11 +5115,8 @@ async def test_process_group_patch_remove_refuses_when_two_members_share_the_id( assert "more than one member of this group" in str(exc_info.value.detail) - @pytest.mark.asyncio -async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else( - mocker, scim_upsert_user_enabled -): +async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else(mocker, scim_upsert_user_enabled): """The canonical user id stays authoritative, including when the same account also holds that value as its email, which is how a SCIM-provisioned account is keyed.""" prisma_client = _member_resolution_prisma( @@ -5147,10 +5178,79 @@ async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_acc assert exc_info.value.status_code == 400 assert "member-id" in str(exc_info.value.detail) create_user_mock.assert_not_called() - assert any( - record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records + assert any(record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_reads_the_exact_id_when_two_other_accounts_fill_the_lookup( + mocker, scim_upsert_user_enabled +): + """A value that is one account's id and two other accounts' identities fills the + bounded lookup with the other two. The account keyed by the value must still be + found, or the id would lose its precedence and a non-canonical type would skip + a member that names a real user.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"shared"}, + teams=set(), + sso_user_id_to_user_id={"shared": "by-sso"}, + email_to_user_id={"shared": "by-email"}, + ) + create_user_mock = mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), ) + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="shared", type="direct")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "shared" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("shared")] + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "shared"}) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_reads_the_user_table_once_per_member(mocker, scim_upsert_user_enabled): + """Every member costs one read of the user table, however it resolves: by its exact + id (which still outranks a non-canonical type), by identity, as a SCIM team, or not + at all. Looking the exact id up on its own before the identity read doubled the + reads of a push, and the identity read is a scan.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"by-id"}, + teams={"by-team"}, + email_to_user_id={"by-email@example.com": "email-user"}, + ) + mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="nobody", key="key")), + ) + + result = await _resolve_group_member_ids( + members=[ + SCIMMember(value="by-id", type="direct"), + SCIMMember(value="by-email@example.com"), + SCIMMember(value="by-team"), + SCIMMember(value="nobody"), + ], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert result.all_member_ids == ["by-id", "email-user", "nobody"] + prisma_client.db.litellm_usertable.find_unique.assert_not_awaited() + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [ + _identity_lookup("by-id"), + _identity_lookup("by-email@example.com"), + _identity_lookup("by-team"), + _identity_lookup("nobody"), + ] @pytest.mark.asyncio @@ -5536,10 +5636,7 @@ async def test_resolve_group_member_ids_admits_member_created_concurrently(mocke the member is still admitted: the id resolves to a real user row, so failing or dropping it would be wrong either way.""" prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set()) - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=[None, LiteLLM_UserTable(user_id="raced-user")] - ) - prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="raced-user")) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", AsyncMock(return_value=None), @@ -5619,3 +5716,196 @@ async def test_patch_group_404s_when_team_deleted_mid_request(mocker): assert exc_info.value.code == "404" assert f"Group not found with ID: {group_id}" in exc_info.value.message + + +_SHADOW_MEMBER_VALUE: Final = "00u1shadow" +_SHADOWED_ACCOUNT: Final = "real-1" +_SHADOWED_GROUP: Final = "grp-eng" + + +def _shadowed_tenant_rows() -> tuple[LiteLLM_UserTable, ...]: + """A placeholder keyed by the raw member value, and the real account that value names by SSO id.""" + return ( + LiteLLM_UserTable(user_id=_SHADOW_MEMBER_VALUE, user_email=_SHADOW_MEMBER_VALUE, teams=[_SHADOWED_GROUP]), + LiteLLM_UserTable(user_id=_SHADOWED_ACCOUNT, user_email="alice@example.com", sso_user_id=_SHADOW_MEMBER_VALUE), + ) + + +def _shadow_tenant_prisma( + mocker: MockerFixture, + *, + rows: Sequence[LiteLLM_UserTable], + keys_owned_by: Mapping[str, int] = MappingProxyType({}), +) -> MagicMock: + """Prisma fake whose user rows are live: deleting one removes it from every later lookup.""" + users: Final[dict[str, LiteLLM_UserTable]] = {row.user_id: row for row in rows} + team: Final = LiteLLM_TeamTable( + team_id=_SHADOWED_GROUP, + members=[_SHADOW_MEMBER_VALUE], + members_with_roles=[Member(user_id=_SHADOW_MEMBER_VALUE, role="user")], + metadata={SCIM_MANAGED_TEAM_METADATA_KEY: True}, + ) + + async def find_unique(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + return users.get(where["user_id"]) + + def clause_matches(row: LiteLLM_UserTable, clause: Mapping[str, object]) -> bool: + if "user_id" in clause: + return row.user_id == clause["user_id"] + if "sso_user_id" in clause: + return row.sso_user_id == clause["sso_user_id"] + email_filter: Final = clause["user_email"] + assert isinstance(email_filter, dict) + return (row.user_email or "").casefold() == str(email_filter["equals"]).casefold() + + async def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: + clauses: Final = where["OR"] + assert isinstance(clauses, list) + matched: Final = tuple(row for row in users.values() if any(clause_matches(row, clause) for clause in clauses)) + return matched[:take] if take else matched + + async def delete(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + return users.pop(where["user_id"], None) + + async def keys_for(where: Mapping[str, object]) -> tuple[MagicMock, ...]: + return tuple(mocker.MagicMock() for _ in range(keys_owned_by.get(str(where["user_id"]), 0))) + + async def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: + return team if where["team_id"] == team.team_id else None + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=find_unique) + prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows) + prisma_client.db.litellm_usertable.delete = AsyncMock(side_effect=delete) + prisma_client.db.litellm_teamtable = mocker.MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team) + prisma_client.db.litellm_verificationtoken = mocker.MagicMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=keys_for) + prisma_client.db.litellm_invitationlink = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + prisma_client.db.litellm_organizationmembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + prisma_client.db.litellm_teammembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + return prisma_client + + +@pytest.fixture +def shadowed_tenant(mocker, monkeypatch, scim_upsert_user_enabled) -> MagicMock: + from litellm.proxy import proxy_server + + prisma_client: Final = _shadow_tenant_prisma(mocker, rows=_shadowed_tenant_rows()) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + return prisma_client + + +async def _push_shadow_member(prisma_client: MagicMock): + return await _resolve_group_member_ids( + members=[SCIMMember(value=_SHADOW_MEMBER_VALUE)], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + +@pytest.mark.asyncio +async def test_merge_placeholder_hands_the_group_to_the_shadowed_account(mocker, shadowed_tenant): + """Every group push of the shadowing value is refused until the placeholder is folded into + the real account; after the merge the same push resolves to that account.""" + team_member_add_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock() + ) + ) + team_member_delete_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock() + ) + ) + + with pytest.raises(HTTPException) as before: + await _push_shadow_member(shadowed_tenant) + assert before.value.status_code == 400 + + result: Final = await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE) + + assert result == SCIMPlaceholderMergeResult( + placeholder_user_id=_SHADOW_MEMBER_VALUE, + merged_into_user_id=_SHADOWED_ACCOUNT, + team_ids=(_SHADOWED_GROUP,), + ) + added: Final = team_member_add_mock.call_args.kwargs["data"] + assert (added.team_id, added.member.user_id) == (_SHADOWED_GROUP, _SHADOWED_ACCOUNT) + dropped: Final = team_member_delete_mock.call_args.kwargs["data"] + assert (dropped.team_id, dropped.user_id) == (_SHADOWED_GROUP, _SHADOW_MEMBER_VALUE) + shadowed_tenant.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"user_id": _SHADOW_MEMBER_VALUE} + ) + shadowed_tenant.db.litellm_usertable.delete.assert_awaited_once_with(where={"user_id": _SHADOW_MEMBER_VALUE}) + + after: Final = await _push_shadow_member(shadowed_tenant) + assert after.all_member_ids == [_SHADOWED_ACCOUNT] + assert after.created_users == [] + + +@pytest.mark.asyncio +async def test_merge_placeholder_keeps_the_placeholder_when_the_roster_write_fails(mocker, shadowed_tenant): + """If the real account cannot join the team, the placeholder stays on it, or the membership is gone + from both accounts.""" + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(side_effect=Exception("database connection lost")), + ) + team_member_delete_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock() + ) + ) + + with pytest.raises(ProxyException): + await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE) + + team_member_delete_mock.assert_not_awaited() + shadowed_tenant.db.litellm_usertable.delete.assert_not_awaited() + assert await shadowed_tenant.db.litellm_usertable.find_unique(where={"user_id": _SHADOW_MEMBER_VALUE}) is not None + + +@pytest.mark.parametrize( + ("rows", "keys_owned_by", "merged", "reason"), + [ + pytest.param(_shadowed_tenant_rows(), {}, _SHADOWED_ACCOUNT, "SSO identity of its own", id="real-account"), + pytest.param( + _shadowed_tenant_rows(), {_SHADOW_MEMBER_VALUE: 2}, _SHADOW_MEMBER_VALUE, "2 virtual keys", id="owns-keys" + ), + pytest.param(_shadowed_tenant_rows()[:1], {}, _SHADOW_MEMBER_VALUE, "shadows no account", id="names-nobody"), + pytest.param( + (*_shadowed_tenant_rows(), LiteLLM_UserTable(user_id="real-2", user_email=_SHADOW_MEMBER_VALUE.upper())), + {}, + _SHADOW_MEMBER_VALUE, + "names 2 accounts (real-1, real-2)", + id="names-two-accounts", + ), + ], +) +@pytest.mark.asyncio +async def test_merge_placeholder_refuses_rows_that_are_not_a_lone_placeholder( + mocker, monkeypatch, scim_upsert_user_enabled, rows, keys_owned_by, merged, reason +): + """Only a row with no SSO identity and no keys whose id names exactly one other account is folded; + anything else could move memberships to the wrong person, so nothing is written.""" + from litellm.proxy import proxy_server + + prisma_client: Final = _shadow_tenant_prisma(mocker, rows=rows, keys_owned_by=keys_owned_by) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + team_member_add_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock() + ) + ) + + with pytest.raises(ProxyException) as exc_info: + await merge_placeholder(user_id=merged) + + assert int(exc_info.value.code) == 409 + assert reason in str(exc_info.value.message) + team_member_add_mock.assert_not_awaited() + prisma_client.db.litellm_usertable.delete.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index a74aa553449..c525af84511 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -881,6 +881,7 @@ def _leg_record(**overrides: object) -> MagicMock: "target_type": "key", "target_id": "key-hash", "router_name": "my-router", + "router_names": (), "direction": "forward", "baseline_model": None, "judge_model": "anthropic/claude-sonnet-5", @@ -931,6 +932,7 @@ def _shadow_prisma( legs=(), agg_rows=None, by_leg_rows=None, + by_router_rows=None, known_keys=("key-hash", "key-hash-2"), key_teams=None, known_teams=None, @@ -1030,6 +1032,7 @@ def _shadow_prisma( "target_type", "target_id", "router_name", + "router_names", "direction", "baseline_model", "judge_model", @@ -1065,6 +1068,8 @@ def _shadow_prisma( return [{"judged_count": 10, "error_count": 2, "judge_spend": 0.031}] if "SELECT job_id AS grp" in sql: return by_leg_rows if by_leg_rows is not None else [] + if "COALESCE(a.router_name" in sql: + return by_router_rows if by_router_rows is not None else [] if 'FROM "LiteLLM_ShadowEvalFunnel"' in sql: return prisma.funnel_rows return agg_rows if agg_rows is not None else [] @@ -1121,7 +1126,15 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once() rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("key", "key-hash-2")] - assert len({frozenset((k, v) for k, v in row.items() if k not in ("target_id", "id")) for row in rows}) == 1 + assert ( + len( + { + frozenset((k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id")) + for row in rows + } + ) + == 1 + ) assert len({row["id"] for row in rows}) == len(rows) assert len({row["group_id"] for row in rows}) == 1 assert all(row["max_turns"] == SHADOW_EVAL_TURN_VALVE and row["created_by"] == "admin" for row in rows) @@ -1138,6 +1151,63 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert all(target.max_turns == SHADOW_EVAL_TURN_VALVE for target in response.targets) +@pytest.mark.asyncio +async def test_start_shadow_eval_multi_router_writes_the_set_on_every_leg(monkeypatch: pytest.MonkeyPatch): + """A multi-router job stores the full set in router_names and the first router in + router_name, so a rolling-deploy pod that predates router_names still runs a valid + single-arm eval and its unstamped attempt rows attribute to that first router.""" + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + response = await start_shadow_eval( + _start_request(router_name=None, router_names=("my-router", "classifier-router")), ADMIN + ) + + rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"] + assert all(row["router_name"] == "my-router" for row in rows) + assert all(row["router_names"] == ["my-router", "classifier-router"] for row in rows) + assert response.router_names == ("my-router", "classifier-router") + assert response.router_name == "my-router" + + +@pytest.mark.asyncio +async def test_start_shadow_eval_rejects_an_unconfigured_router_in_the_set(monkeypatch: pytest.MonkeyPatch): + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException, match="not-a-router") as exc: + await start_shadow_eval(_start_request(router_name=None, router_names=("my-router", "not-a-router")), ADMIN) + + assert exc.value.status_code == 400 + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_judge_collision_is_found_on_every_router_of_the_set(monkeypatch: pytest.MonkeyPatch): + """The judge-as-candidate guard walks every candidate router: a judge that serves an + arm of the SECOND router still poisons the whole job's win rates.""" + import litellm.proxy.proxy_server as proxy_server + + _configure_anthropic_sdk_judge(monkeypatch) + prisma = _shadow_prisma() + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException, match="also an arm") as exc: + await start_shadow_eval(_start_request(router_name=None, router_names=("my-router", "sonnet-router")), ADMIN) + + assert exc.value.status_code == 400 + prisma.db.litellm_shadowevaljob.create_many.assert_not_called() + + @pytest.mark.asyncio async def test_start_shadow_eval_rejects_an_uncredentialed_sdk_judge(monkeypatch: pytest.MonkeyPatch) -> None: import litellm @@ -1798,6 +1868,63 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke assert error_where == {"job_id": {"in": ["leg-1", "leg-2"]}, "outcome": "error"} +@pytest.mark.asyncio +async def test_get_shadow_eval_job_slices_results_per_router(monkeypatch: pytest.MonkeyPatch): + """A multi-router job's detail carries one slice per arm, aggregated by the arm + stamped on each attempt row, with unstamped legacy rows attributed to the job's own + router by the read (the COALESCE against the leg's router_name).""" + import litellm.proxy.proxy_server as proxy_server + + def agg(grp: str, wins: int) -> dict[str, object]: + return { + "grp": grp, + "turn_count": 4, + "real_wins": 4 - wins, + "shadow_wins": wins, + "ties": 0, + "avg_confidence": 0.8, + "real_spend": 0.08, + "shadow_spend": 0.02, + "cache_hit_turns": 0, + } + + prisma = _shadow_prisma( + legs=[_leg_record(router_names=("my-router", "alt-router"))], + agg_rows=[agg("SIMPLE", 3)], + by_router_rows=[agg("my-router", 1), agg("alt-router", 3)], + ) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response = await get_shadow_eval_job("job-1", VIEWER) + + assert response.router_names == ("my-router", "alt-router") + assert response.router_name == "my-router" + assert [(s.group, s.shadow_win_rate_pct) for s in response.results.by_router] == [ + ("my-router", 25.0), + ("alt-router", 75.0), + ] + router_sql = next( + call.args[0] for call in prisma.db.query_raw.await_args_list if "COALESCE(a.router_name" in call.args[0] + ) + assert "COALESCE(a.router_name, j.router_name)" in router_sql + assert 'JOIN "LiteLLM_ShadowEvalJob" j ON j.id = a.job_id' in router_sql + assert "a.job_id = ANY($1::text[])" in router_sql + + +@pytest.mark.asyncio +async def test_job_responses_resolve_router_names_with_legacy_fallback(monkeypatch: pytest.MonkeyPatch): + """Rows from before router_names existed carry their whole set in router_name.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(legs=[_leg_record(router_names=())]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + response = await get_shadow_eval_job("job-1", VIEWER) + + assert response.router_names == ("my-router",) + assert response.router_name == "my-router" + + @pytest.mark.asyncio async def test_get_shadow_eval_job_404s_and_gates_on_role(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 1c80aa5683f..7e2e680743f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11033,6 +11033,123 @@ class TestLIT1884KeyUpdateValidation: ) +class TestLIT4891SafePresetKeyTypeTransition: + def _make_existing_key(self, allowed_routes): + row = MagicMock() + row.user_id = "internal-user-123" + row.created_by = "internal-user-123" + row.token = "hashed_token" + row.team_id = None + row.max_budget = None + row.spend = 0.0 + row.organization_id = None + row.project_id = None + row.allowed_routes = allowed_routes + return row + + def _make_auth(self): + return UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + async def _run_update(self, data, existing_key_row): + try: + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=self._make_auth(), + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + except HTTPException as exc: + return exc + return None + + def _assert_routes_403(self, exc): + assert exc is not None + assert exc.status_code == 403 + assert "Only proxy admins can set" in str(exc.detail) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_clear_safe_preset_to_full_access(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_switch_full_access_to_safe_preset(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=[]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_owner_can_narrow_to_read_only_preset(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_can_resend_read_only_preset_unchanged(self): + assert ( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["info_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) + is None + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_widen_read_only_key_to_full_access(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_widen_read_only_key_to_llm_api(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["llm_api_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["info_routes"]), + ) + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_clear_custom_route_restriction(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=[]), + existing_key_row=self._make_existing_key(allowed_routes=["/chat/completions"]), + ) + ) + + @pytest.mark.asyncio + async def test_non_admin_cannot_set_non_preset_routes(self): + self._assert_routes_403( + await self._run_update( + data=UpdateKeyRequest(key="sk-test", allowed_routes=["management_routes"]), + existing_key_row=self._make_existing_key(allowed_routes=["llm_api_routes"]), + ) + ) + + class TestKeyOwnerPrivilegeEscalation: """ Policy: @@ -12007,9 +12124,10 @@ class TestAllowedRoutesCallerPermission: @pytest.mark.asyncio async def test_non_admin_update_key_explicit_empty_allowed_routes_rejected(self): - """`update_key_fn` rejects a non-admin when `allowed_routes` is - present as `[]` in the request body. The value matches the model - default but `model_fields_set` distinguishes the two.""" + """`update_key_fn` rejects a non-admin clearing a custom (non-preset) + route restriction with an explicit `[]` in the request body. The value + matches the model default but `model_fields_set` distinguishes the + two. Clearing from a safe preset is allowed (LIT-4891).""" from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, ) @@ -12032,7 +12150,7 @@ class TestAllowedRoutesCallerPermission: patch( "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=MagicMock(allowed_routes=["/chat/completions"]), ), ): with pytest.raises(ProxyException) as exc_info: @@ -12047,8 +12165,8 @@ class TestAllowedRoutesCallerPermission: @pytest.mark.asyncio async def test_non_admin_update_key_explicit_null_allowed_routes_rejected(self): - """`update_key_fn` rejects a non-admin when `allowed_routes` is - present as `null` in the request body.""" + """`update_key_fn` rejects a non-admin clearing a custom (non-preset) + route restriction with an explicit `null` in the request body.""" from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, ) @@ -12071,7 +12189,7 @@ class TestAllowedRoutesCallerPermission: patch( "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key", new_callable=AsyncMock, - return_value=MagicMock(), + return_value=MagicMock(allowed_routes=["/chat/completions"]), ), ): with pytest.raises(ProxyException) as exc_info: @@ -14142,6 +14260,311 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): mock_prisma_client.db.query_raw.assert_not_awaited() +@pytest.mark.asyncio +async def test_info_key_fn_reports_budget_limits_usage(monkeypatch): + """ + /key/info reports current-window spend per budget window under budget_limits_usage, + keyed by budget_duration and read from the same counter enforcement uses, while + budget_limits itself comes back exactly as stored. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + } + ] + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.73) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-w" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [dict(w) for w in budget_limits], + "user_id": "user-w", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-window-key", + ) + + result = await info_key_fn( + key="sk-test-window-key", + user_api_key_dict=user_api_key_dict, + ) + + assert result["info"]["budget_limits"] == budget_limits + assert result["info"]["budget_limits_usage"] == {"1h": {"current_spend": 0.73}} + + mock_get_current_spend.assert_awaited_once() + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == f"spend:key:{test_key_token}:window:1h" + assert call_kwargs["max_budget"] == 2.0 + assert call_kwargs["window_entity_type"] == "Key" + assert call_kwargs["window_entity_id"] == test_key_token + assert call_kwargs["window_duration"] == "1h" + assert call_kwargs["window_start"] is not None + + +@pytest.mark.asyncio +async def test_info_key_fn_no_budget_limits_skips_spend_lookup(monkeypatch): + """Keys without budget windows get no budget_limits_usage field and trigger no spend lookup.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + test_key_token = "hashed_token_no_windows" + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=0.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key_info = MagicMock(spec=LiteLLM_VerificationToken) + mock_key_info.token = test_key_token + mock_key_info.object_permission_id = None + mock_key_info.user_id = "user-nw" + mock_key_info.team_id = None + mock_key_info.litellm_budget_table = None + mock_key_info.model_dump.return_value = { + "token": test_key_token, + "budget_limits": None, + "user_id": "user-nw", + "team_id": None, + "object_permission_id": None, + "litellm_budget_table": None, + } + mock_key_info.dict.return_value = mock_key_info.model_dump.return_value + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_info + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-test-no-window-key", + ) + + result = await info_key_fn( + key="sk-test-no-window-key", + user_api_key_dict=user_api_key_dict, + ) + + assert result["info"]["budget_limits"] is None + assert "budget_limits_usage" not in result["info"] + mock_get_current_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_info_key_fn_v2_reports_budget_limits_usage(monkeypatch): + """/v2/key/info reports budget_limits_usage per window and leaves budget_limits as stored.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken + from litellm.proxy.management_endpoints.key_management_endpoints import ( + info_key_fn_v2, + ) + + test_key_token = "hashed_token_v2_window_test" + budget_limits = [ + { + "reset_at": "2026-08-15T18:00:00+00:00", + "max_budget": 2.0, + "budget_duration": "1h", + }, + { + "reset_at": "2026-08-16T00:00:00+00:00", + "max_budget": 20.0, + "budget_duration": "1d", + }, + ] + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_user_api_key_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + mock_get_current_spend = AsyncMock(return_value=1.25) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_key_token + mock_key.user_id = "user-v2-w" + mock_key.team_id = None + mock_key.model_dump.return_value = { + "token": test_key_token, + "budget_limits": [dict(w) for w in budget_limits], + "user_id": "user-v2-w", + "team_id": None, + "litellm_budget_table": None, + } + mock_key.dict.return_value = mock_key.model_dump.return_value + + mock_prisma_client.get_data = AsyncMock(return_value=[mock_key]) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin-v2-w", + ) + + result = await info_key_fn_v2( + data=KeyRequest(keys=[test_key_token]), + user_api_key_dict=user_api_key_dict, + ) + + assert len(result["info"]) == 1 + assert result["info"][0]["budget_limits"] == budget_limits + assert result["info"][0]["budget_limits_usage"] == { + "1h": {"current_spend": 1.25}, + "1d": {"current_spend": 1.25}, + } + assert mock_get_current_spend.await_count == 2 + counter_keys = { + call.kwargs["counter_key"] for call in mock_get_current_spend.await_args_list + } + assert counter_keys == { + f"spend:key:{test_key_token}:window:1h", + f"spend:key:{test_key_token}:window:1d", + } + assert { + call.kwargs["window_duration"] for call in mock_get_current_spend.await_args_list + } == {"1h", "1d"} + + +@pytest.mark.asyncio +async def test_build_budget_limits_usage_json_string_input(monkeypatch): + """budget_limits stored as a JSON string is parsed and reported per window.""" + import json as json_module + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.5) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + raw = json_module.dumps( + [{"budget_duration": "1h", "max_budget": 2.0, "reset_at": None}] + ) + result = await _build_budget_limits_usage(budget_limits=raw, api_key_hash="hash-1") + + assert result == {"1h": {"current_spend": 0.5}} + mock_get_current_spend.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_build_budget_limits_usage_empty_windows_returns_none(monkeypatch): + """A key with no windows (None, [], or "[]") returns None so the field is left off; no spend lookup runs.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + for stored in (None, [], "[]"): + assert await _build_budget_limits_usage(budget_limits=stored, api_key_hash="hash-1") is None + mock_get_current_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_build_budget_limits_usage_window_without_max_budget(monkeypatch): + """A window with only budget_duration still reports current_spend, read without a budget ceiling.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=0.75) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + result = await _build_budget_limits_usage( + budget_limits=[{"budget_duration": "2d"}], api_key_hash="hash-no-max" + ) + + assert result == {"2d": {"current_spend": 0.75}} + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == "spend:key:hash-no-max:window:2d" + assert call_kwargs["window_duration"] == "2d" + assert call_kwargs["max_budget"] is None + + +@pytest.mark.asyncio +async def test_build_budget_limits_usage_pydantic_windows(monkeypatch): + """BudgetLimitEntry windows (the shape UserAPIKeyAuth carries) are dumped to dicts and reported.""" + from unittest.mock import AsyncMock + + from litellm.models.team import BudgetLimitEntry + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_budget_limits_usage, + ) + + mock_get_current_spend = AsyncMock(return_value=1.0) + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ) + + result = await _build_budget_limits_usage( + budget_limits=[BudgetLimitEntry(budget_duration="7d", max_budget=10.0)], + api_key_hash="hash-2", + ) + + assert result == {"7d": {"current_spend": 1.0}} + call_kwargs = mock_get_current_spend.await_args.kwargs + assert call_kwargs["counter_key"] == "spend:key:hash-2:window:7d" + assert call_kwargs["window_duration"] == "7d" + assert call_kwargs["max_budget"] == 10.0 + + @pytest.mark.asyncio async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): """/key/info reads the one counter enforcement reads: the configured budget model. @@ -17066,3 +17489,49 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project assert exc_info.value.status_code == 400 assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] + + +def test_generate_key_request_blank_team_id_is_personal(): + """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925).""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_team_key, + ) + + cleared = GenerateKeyRequest(team_id="") + assert cleared.team_id is None + assert _is_team_key(data=cleared) is False + assert RegenerateKeyRequest(team_id="").team_id is None + assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" + + +def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): + """key_generation_check with team_id="" must take the personal-key path instead + of failing the team lookup with "Unable to find team object" (LIT-3925).""" + from litellm.proxy._types import KeyManagementRoutes + from litellm.proxy.management_endpoints.key_management_endpoints import ( + key_generation_check, + ) + + monkeypatch.setattr( + litellm, + "key_generation_settings", + { + "team_key_generation": {"allowed_team_member_roles": ["admin"]}, + "personal_key_generation": {"allowed_user_roles": ["proxy_admin", "internal_user"]}, + }, + ) + + assert ( + key_generation_check( + team_table=None, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + data=GenerateKeyRequest(key_alias="personal", team_id=""), + route=KeyManagementRoutes.KEY_GENERATE, + ) + is True + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 393953ccf68..4661cc17dbc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -18,6 +18,7 @@ from litellm.proxy._types import ( ReconcileOutcome, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, _get_team_deployments, @@ -263,6 +264,131 @@ class TestModelManagementAuthChecks: ) assert "403" in str(exc_info.value) + def test_can_user_attach_credential_admin_success(self): + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + user_api_key_dict=self.admin_user, + ) + assert result is True + + def test_can_user_attach_credential_without_credential_allows_any_role(self): + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model"), + user_api_key_dict=self.team_admin_user, + ) + assert result is True + + def test_can_user_attach_credential_team_admin_fails(self): + with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info: + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + + def test_can_user_attach_credential_unchanged_existing_allows_any_role(self): + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + ) + assert result is True + + def test_can_user_attach_credential_unchanged_encrypted_existing_allows_any_role(self, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + encrypted_name = encrypt_value_helper(value="shared-credential") + assert encrypted_name != "shared-credential" + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name=encrypted_name), + ) + assert result is True + + @pytest.mark.asyncio + async def test_add_new_model_rejects_credential_attach_for_non_admin(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + mock_prisma = MagicMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the credential check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="credential-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", litellm_credential_name="shared-credential" + ), + model_info={"id": "credential-create-test"}, + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + mock_prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_patch_model_rejects_credential_attach_for_non_admin(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "credential-patch-test" + db_model = Deployment( + model_name="credential-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info={"id": model_id}, + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: stubs the DB row fetch; only the credential check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=db_model), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the credential check is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: asserts the DB write is never reached on rejection + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(), + ) as mock_update, + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams( + model="openai/gpt-4o", litellm_credential_name="shared-credential" + ) + ), + user_api_key_dict=self.team_admin_user, + ) + assert exc_info.value.code == "403" + mock_update.assert_not_awaited() + + def test_can_user_attach_credential_internal_user_fails(self): + with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info: + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), + user_api_key_dict=self.normal_user, + ) + assert exc_info.value.code == "403" + class MockModelTable: def __init__(self, model_aliases: Dict[str, str], include: Optional[dict] = None): diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ffa6bc601e9..30b2ab86b9a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3741,6 +3741,93 @@ async def test_list_team_v2_with_status_deleted(): assert len(result["teams"]) == 2 +@pytest.mark.asyncio +async def test_list_team_v2_includes_litellm_model_table(): + """ + Regression test for GH #26312: GET /v2/team/list must eagerly load the + litellm_model_table relation for active teams, same as /team/info and + /team/list, or a team's model_aliases always read back as null from this + endpoint. Deleted teams are excluded: LiteLLM_DeletedTeamTable has no such + relation in the Prisma schema, so requesting it there raises + UnknownRelationalFieldError against a real database. + + The fake find_many below only attaches litellm_model_table when its own + `include` kwarg actually asks for the relation, so the assertions below + are on what the caller gets back, not on how find_many was called. + """ + from unittest.mock import AsyncMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 + + mock_request = Mock(spec=Request) + mock_user_api_key_dict_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user_123", + ) + + def _team_row(team_id: str, include) -> Mock: + model_table = ( + { + "id": 1, + "model_aliases": {"my-fast-model": "fake-model"}, + "created_by": "u", + "updated_by": "u", + "team": None, + } + if (include or {}).get("litellm_model_table") + else None + ) + return Mock( + team_id=team_id, + model_dump=lambda: { + "team_id": team_id, + "team_alias": "t", + "litellm_model_table": model_table, + }, + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: # test-quality-ok: this file's DB-mock convention + mock_db = Mock() + mock_prisma_client.db = mock_db + + mock_db.litellm_teamtable.find_many = AsyncMock( + side_effect=lambda **kw: [_team_row("team_1", kw.get("include"))] + ) + mock_db.litellm_teamtable.count = AsyncMock(return_value=1) + mock_db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) + + result = await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status=None, + ) + + assert result["teams"][0].litellm_model_table is not None + assert result["teams"][0].litellm_model_table.model_aliases == {"my-fast-model": "fake-model"} + + mock_db.litellm_deletedteamtable.find_many = AsyncMock( + side_effect=lambda **kw: [_team_row("team_2", kw.get("include"))] + ) + mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=1) + + await list_team_v2( + http_request=mock_request, + user_id=None, + user_api_key_dict=mock_user_api_key_dict_admin, + page=1, + page_size=10, + status="deleted", + ) + + assert "include" not in mock_db.litellm_deletedteamtable.find_many.call_args.kwargs + + @pytest.mark.asyncio async def test_list_team_v2_org_admin_sees_org_teams(): """ diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py new file mode 100644 index 00000000000..60c36e33e09 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.management_helpers.access_group_key_sync import ( + sync_key_access_group_membership, + sync_key_regeneration_access_group_membership, +) + + +def _routed_prisma_client(): + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.query_raw = AsyncMock(return_value=[]) + reader_inner.query_raw = AsyncMock(return_value=[]) + writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) + reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + return SimpleNamespace(db=routing), writer_inner, reader_inner + + +@pytest.mark.asyncio +async def test_regeneration_repoint_update_runs_on_the_writer(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client() + + await sync_key_regeneration_access_group_membership( + prisma_client=prisma_client, + previous_key_token="old-token", + new_key_token="new-token", + data=None, + existing_key_row=MagicMock(), + ) + + writer_inner.query_raw.assert_awaited_once() + assert writer_inner.query_raw.await_args.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') + reader_inner.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_membership_attach_and_detach_updates_run_on_the_writer(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client() + + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token="token", + previous_access_group_ids=["ag-old"], + updated_access_group_ids=["ag-new"], + ) + + assert writer_inner.query_raw.await_count == 2 + assert all( + call.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') for call in writer_inner.query_raw.await_args_list + ) + reader_inner.query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 5ef83344c1a..f2b6b799271 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -1,11 +1,9 @@ import json +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException - -from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import ( LiteLLM_ObjectPermissionBase, LiteLLM_ObjectPermissionTable, @@ -13,10 +11,10 @@ from litellm.proxy._types import ( SpecialMCPServerName, ) from litellm.proxy.management_helpers.object_permission_utils import ( + _drop_stale_object_permission_mcp_servers, _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, _resolve_team_allowed_mcp_servers, - _rewrite_object_permission_mcp_servers, _set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, validate_key_mcp_servers_against_team, @@ -153,10 +151,10 @@ def test_extract_requested_mcp_server_ids_excludes_no_mcp_servers_sentinel(): assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"} -def test_rewrite_object_permission_mcp_servers_preserves_sentinel(): - obj_perm = {"mcp_servers": ["no-mcp-servers", "alias-1"]} - _rewrite_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}}) - assert obj_perm["mcp_servers"] == ["no-mcp-servers", "server-1"] +def test_drop_stale_object_permission_mcp_servers_preserves_sentinel_and_alias(): + obj_perm = {"mcp_servers": ["no-mcp-servers", "alias-1", "gone-id"]} + _drop_stale_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}, "gone-id": set()}) + assert obj_perm["mcp_servers"] == ["no-mcp-servers", "alias-1"] @pytest.mark.asyncio @@ -692,9 +690,10 @@ async def test_validate_mcp_server_alias_outside_team_scope_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_mcp_server_alias_is_normalized_before_save( - mock_access_groups, mock_allow_all -): +async def test_validate_mcp_server_alias_persists_verbatim(mock_access_groups, mock_allow_all): + """Regression for the multi-region shared-DB setup: an alias grant must be + stored as the alias, so every instance can expand it to its own local id. + Rewriting to this instance's server_id breaks access on the other region.""" team_obj = _make_team_obj(mcp_servers=["allowed-server-id"]) object_permission = { "mcp_servers": ["allowed-alias"], @@ -706,8 +705,27 @@ async def test_validate_mcp_server_alias_is_normalized_before_save( team_obj=team_obj, ) - assert object_permission["mcp_servers"] == ["allowed-server-id"] - assert object_permission["mcp_tool_permissions"] == {"allowed-server-id": ["tool1"]} + assert object_permission["mcp_servers"] == ["allowed-alias"] + assert object_permission["mcp_tool_permissions"] == {"Allowed Server": ["tool1"]} + + +def test_alias_grant_expands_on_other_region_after_save(): + """Cross-region flow: the west instance saves an alias grant (its resolver maps + the alias to west's hash-derived id), then the central instance, whose registry + maps the same alias to a different id, expands the persisted grant. Rewriting + to west's id at save time is exactly the regression this guards against.""" + west_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("west-id", alias="github-mcp")]) + central_mgr = _make_mock_mcp_manager(servers=[_make_mock_mcp_server("central-id", alias="github-mcp")]) + + object_permission = {"mcp_servers": ["github-mcp"]} + _drop_stale_object_permission_mcp_servers(object_permission, {"github-mcp": {"west-id"}}) + assert object_permission["mcp_servers"] == ["github-mcp"] + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + expand = MCPServerManager.expand_permission_list + assert expand(west_mgr, object_permission["mcp_servers"]) == ["west-id"] + assert expand(central_mgr, object_permission["mcp_servers"]) == ["central-id"] @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f5ae0fe5977..d3f17c73499 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2,6 +2,7 @@ import asyncio import json import logging import os +from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO from types import SimpleNamespace @@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( websocket_passthrough_request, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, @@ -5464,7 +5466,10 @@ def test_the_marker_check_distinguishes_the_two_route_kinds(): assert request_dispatched_to_pass_through_endpoint(builtin) is False -async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: UserAPIKeyAuth) -> tuple[int, object]: +async def _drive_passthrough_request_and_capture_logging( + user_api_key_dict: UserAPIKeyAuth, + on_pre_call: Callable[[LiteLLMLoggingObj | None], None] | None = None, +) -> tuple[int, LiteLLMLoggingObj | None]: import litellm from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -5487,10 +5492,12 @@ async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: User mock_request.query_params = QueryParams({}) mock_request.body = AsyncMock(return_value=b'{"model": "gemini-2.0-flash"}') - captured_data: dict = {} + captured_data: dict = {} # mutable-ok: the pre-call hook records the request data into it async def capture_pre_call_hook(user_api_key_dict, data, call_type): captured_data.update(data) + if on_pre_call is not None: + on_pre_call(data.get("litellm_logging_obj")) return data mock_proxy_logging = MagicMock() @@ -5623,3 +5630,103 @@ async def test_resolve_team_callback_wiring_fails_open_on_operational_error(): assert wiring.success_callbacks is None assert wiring.failure_callbacks is None assert wiring.logging_kwargs is None + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_guardrail_readable_metadata(): + """A pre-call guardrail reads the request headers off the passthrough logging + params without raising.""" + from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import ( + _logged_request_headers, + ) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success_and_failure", + "callback_vars": { + "langfuse_public_key": "pk_test", + "langfuse_secret_key": "sk_test", + }, + } + ] + }, + ) + + observed: dict[str, dict[str, str] | BaseException] = {} # mutable-ok: the pre-call hook records into it + + def read_headers_the_way_a_guardrail_does(logging_obj: LiteLLMLoggingObj | None) -> None: + assert logging_obj is not None + try: + observed["headers"] = _logged_request_headers(logging_obj) + except Exception as exc: # noqa: BLE001 - the regression is that this used to raise + observed["headers"] = exc + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging( + user_api_key_dict, on_pre_call=read_headers_the_way_a_guardrail_does + ) + + assert "headers" in observed, "the pre-call hook never ran, so nothing was observed" + assert observed["headers"] == {}, f"guardrail header read failed: {observed['headers']!r}" + assert status_code == 200 + assert logging_obj is not None + assert logging_obj.dynamic_success_callbacks, "team success callbacks must stay wired" + assert logging_obj.standard_callback_dynamic_params.get("langfuse_public_key") == "pk_test" + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_cost_router_logger_working(): + """The cost router's logger reads the deployment id off the passthrough logging + params without raising. least_busy shares the read but swallows the exception, + so this is the strategy where the break is observable.""" + from litellm._logging import verbose_logger + from litellm.caching.caching import DualCache + from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler + + handler = LowestCostLoggingHandler(router_cache=DualCache()) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success_and_failure", + "callback_vars": { + "langfuse_public_key": "pk_test", + "langfuse_secret_key": "sk_test", + }, + } + ] + }, + ) + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict) + assert status_code == 200 + assert logging_obj is not None + + raised: list[logging.LogRecord] = [] # mutable-ok: logging.Handler records into it + + class _RecordTracebacks(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.exc_info is not None: + raised.append(record) + + recorder = _RecordTracebacks() + verbose_logger.addHandler(recorder) + try: + await handler.async_log_success_event( + kwargs=logging_obj.model_call_details, + response_obj=None, + start_time=None, + end_time=None, + ) + finally: + verbose_logger.removeHandler(recorder) + + assert not raised, f"cost router logger raised on the passthrough logging params: {raised[0].exc_info}" diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 054a5af4148..4fcb7d22588 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -749,6 +749,106 @@ async def test_single_step_pipeline_allow(monkeypatch): assert guard.calls == 1 +@pytest.mark.asyncio +async def test_allow_restores_independent_guardrails_list(monkeypatch): + """ + Request activates an independent guardrail; an unrelated pipeline runs and allows. + Expected: no modified_data escapes, so the request's guardrails list survives + and the independent guardrail still runs at later lifecycle stages (post_call). + Regression: LIT-6587 (pipeline clobbered the list with its last step's guardrail). + """ + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = { + "messages": [{"role": "user", "content": "clean content"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert pipeline_guard.calls == 1 + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert propagated["metadata"]["guardrails"] == ["independent-output-guard"] + assert data["metadata"]["guardrails"] == ["independent-output-guard"] + + +@pytest.mark.asyncio +async def test_allow_does_not_leak_guardrails_into_bare_request(monkeypatch): + """A request without metadata must not gain a metadata.guardrails list from the pipeline.""" + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = {"messages": [{"role": "user", "content": "clean content"}]} + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert "guardrails" not in propagated.get("metadata", {}) + assert "metadata" not in data + + +@pytest.mark.asyncio +async def test_data_forwarding_keeps_changes_and_restores_guardrails_list(monkeypatch): + """A pass_data pipeline's modifications propagate while the request's guardrails list is restored.""" + pii_guard = PiiMaskingGuardrail(guardrail_name="pii-masker") + content_guard = ContentCheckGuardrail(guardrail_name="content-check") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="pii-masker", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="content-check", on_fail="block", on_pass="allow"), + ], + ) + + monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard]) + + data = { + "messages": [{"role": "user", "content": "Hello John Smith"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="pii-then-safety", + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" + assert result.modified_data["metadata"]["guardrails"] == ["independent-output-guard"] + + @pytest.mark.asyncio async def test_step_results_include_duration(monkeypatch): """Step results should include timing information.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index 74542a3eaf6..de76c7257cf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -12,13 +12,16 @@ from __future__ import annotations import io from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from litellm.proxy import proxy_server +from litellm.types.llms.openai import HttpxBinaryResponseContent @pytest.fixture -def patched_speech(monkeypatch): +def patched_speech(monkeypatch, request): + upstream_content_type = getattr(request, "param", "audio/mpeg") monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) monkeypatch.setattr( proxy_server, @@ -36,15 +39,14 @@ def patched_speech(monkeypatch): monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) - class _FakeBinaryResp: - async def aiter_bytes(self, chunk_size: int = 8192): - async def _gen(): - yield b"\x00\x01\x02" - - return _gen() - async def _llm_call(): - return _FakeBinaryResp() + return HttpxBinaryResponseContent( + httpx.Response( + status_code=200, + headers={} if upstream_content_type is None else {"content-type": upstream_content_type}, + content=b"\x00\x01\x02", + ) + ) async def _fake_route_request(*args, **kwargs): return _llm_call() @@ -79,6 +81,24 @@ def patched_speech_error(monkeypatch): yield +@pytest.fixture +def patched_speech_provider_rejection(monkeypatch, patched_speech_error): + import litellm + + async def _raise(*args, **kwargs): + raise litellm.BadRequestError( + message=( + "Gemini TTS only produces raw PCM16 audio, so response_format='mp3' is not supported." + " Supported response formats: pcm, wav." + ), + model="gemini-3.1-flash-tts-preview", + llm_provider="gemini", + ) + + monkeypatch.setattr(proxy_server, "route_request", _raise) + yield + + @pytest.fixture def patched_transcription(monkeypatch): router = MagicMock() @@ -152,6 +172,35 @@ def test_audio_speech_happy_path(client, auth_as, patched_speech, path): } +@pytest.mark.parametrize( + ("patched_speech", "response_format", "expected_content_type"), + [ + ("audio/wav", "wav", "audio/wav"), + ("audio/flac", "flac", "audio/flac"), + ("audio/pcm", "pcm", "audio/pcm"), + ("audio/wav", "mp3", "audio/wav"), + ("application/json", "flac", "audio/flac"), + (None, "wav", "audio/wav"), + (None, None, "audio/mpeg"), + ], + indirect=["patched_speech"], +) +def test_audio_speech_content_type_matches_audio_format( + client, auth_as, patched_speech, response_format, expected_content_type +): + """Regression for LIT-6482: /v1/audio/speech mislabeled wav/flac/pcm as audio/mpeg.""" + payload = { + "model": "tts-1", + "input": "Hi", + "voice": "alloy", + **({} if response_format is None else {"response_format": response_format}), + } + with auth_as(): + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 200 + assert response.headers.get("content-type", "").split(";")[0] == expected_content_type + + @pytest.mark.parametrize("path", ["/v1/audio/speech", "/audio/speech"]) def test_audio_speech_error(client, auth_as, patched_speech_error, path): """Pins ``POST /v1/audio/speech`` and ``POST /audio/speech`` (error).""" @@ -162,6 +211,18 @@ def test_audio_speech_error(client, auth_as, patched_speech_error, path): assert len(response.content) > 0 +def test_audio_speech_bad_request_maps_to_400(client, auth_as, patched_speech_provider_rejection): + """Regression for LIT-6501: a BadRequestError from the speech path surfaced as a generic 500.""" + payload = {"model": "gemini-tts", "input": "Hi", "voice": "Kore", "response_format": "mp3"} + with auth_as(): + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 400 + error = response.json()["error"] + assert "response_format='mp3'" in error["message"] + assert "pcm" in error["message"] + assert "wav" in error["message"] + + @pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"]) def test_audio_transcription_happy_path(client, auth_as, patched_transcription, path): """Pins ``POST /v1/audio/transcriptions`` / ``POST /audio/transcriptions`` (happy).""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index 2b126b1ea95..bc6106a06f8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -45,6 +45,7 @@ def patched_models(monkeypatch): deployment = MagicMock() deployment.litellm_params.model = "gpt-4" router.get_deployment_by_model_group_name = MagicMock(return_value=deployment) + router.get_configured_display_name = MagicMock(return_value=None) monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -187,6 +188,83 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as assert (claude["max_input_tokens"], claude["max_tokens"]) == (500000, 4096) +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_uses_configured_display_name(client, auth_as, patched_models, path): + """A deployment's ``model_info.display_name`` becomes the Anthropic-native + ``display_name`` so Claude Code's picker shows a clean name while the id keeps + routing; models without one keep the id fallback, and the OpenAI-shaped + listing carries no display_name either way.""" + + def _configured(model_name): + return "Kimi K3" if model_name == "gpt-4" else None + + patched_models.get_configured_display_name = MagicMock(side_effect=_configured) + + with auth_as(): + anthropic_response = client.get(path, headers={"anthropic-version": "2023-06-01"}) + openai_response = client.get(path) + + assert anthropic_response.status_code == 200 + gpt_4, claude = anthropic_response.json()["data"] + assert (gpt_4["id"], gpt_4["display_name"]) == ("gpt-4", "Kimi K3") + assert (claude["id"], claude["display_name"]) == ("claude-sonnet", "claude-sonnet") + + assert openai_response.status_code == 200 + openai_models = openai_response.json()["data"] + assert [m["id"] for m in openai_models] == ["gpt-4", "claude-sonnet"] + assert all("display_name" not in m for m in openai_models) + + +@pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) +def test_anthropic_display_name_resolved_via_internal_team_key( + client, auth_as, patched_models, monkeypatch, params +): + """For a team-scoped row the configured display name must be looked up by the + internal routing key while the entry itself is keyed by the public name, so + the clean name lands on the id the client actually sees.""" + from litellm.proxy import utils as proxy_utils + from litellm.proxy.auth import model_checks + + internal_name = "model_name_team-1_c0ffee" + + patched_models.get_model_list = MagicMock( + return_value=[ + { + "model_name": internal_name, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "gpt-4-team", + }, + } + ] + ) + patched_models.get_model_names = MagicMock(return_value=[internal_name]) + patched_models.get_configured_display_name = MagicMock( + side_effect=lambda model_name: "Team GPT" if model_name == internal_name else None + ) + + async def _fake_get_available_models_for_user(**kwargs): + return [internal_name] + + monkeypatch.setattr( + proxy_utils, + "get_available_models_for_user", + _fake_get_available_models_for_user, + ) + monkeypatch.setattr( + model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] + ) + + with auth_as(): + response = client.get( + "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} + ) + + assert response.status_code == 200 + (entry,) = response.json()["data"] + assert (entry["id"], entry["display_name"]) == ("gpt-4-team", "Team GPT") + + @pytest.mark.parametrize("path", ["/v1/models", "/models"]) def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path): """Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope).""" diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index 35b5c72f92e..1e1436fcef8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -9,12 +9,13 @@ Pins (PR2): from __future__ import annotations import asyncio -from unittest.mock import AsyncMock, MagicMock +import json import pytest import litellm from litellm.proxy import proxy_server +from litellm.router_utils import pattern_match_deployments from .conftest import normalize # type: ignore[import-not-found] @@ -99,6 +100,7 @@ def test_token_counter_missing_input_returns_400( @pytest.fixture def patched_supported_params(monkeypatch): + monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr( litellm, "get_llm_provider", @@ -124,12 +126,104 @@ def test_supported_openai_params_happy_path(client, auth_as, patched_supported_p } +def test_supported_openai_params_resolves_router_alias(client, auth_as, monkeypatch): + """A router alias absent from the cost map resolves through the deployment's underlying model.""" + router = litellm.Router( + model_list=[ + { + "model_name": "claude-opus-4-6-cached", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with auth_as(): + response = client.get("/utils/supported_openai_params", params={"model": "claude-opus-4-6-cached"}) + + assert response.status_code == 200 + expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic") + assert response.json() == {"supported_openai_params": expected} + assert "max_tokens" in response.json()["supported_openai_params"] + + +def test_supported_openai_params_declared_prefix_alias_resolves_through_router(client, auth_as, monkeypatch): + """Regression: an alias whose name starts with an authenticating provider's prefix skipped + router resolution and answered with that provider's params instead of the deployment's.""" + router = litellm.Router( + model_list=[ + { + "model_name": "github_copilot/gpt-4o", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with auth_as(): + response = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"}) + + assert response.status_code == 200 + expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic") + assert response.json() == {"supported_openai_params": expected} + + +def test_supported_openai_params_never_runs_oauth_for_authenticating_providers(client, auth_as, monkeypatch, tmp_path): + """Regression: github_copilot/chatgpt names answer from their declaration; resolving them + through ``get_llm_provider`` would run the provider's OAuth device flow and block the event loop.""" + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "access-token").write_text("fake-access-token") + (tmp_path / "api-key.json").write_text( + json.dumps( + { + "token": "fake-api-key", + "expires_at": 4102444800, + "endpoints": {"api": "https://api.githubcopilot.com"}, + } + ) + ) + router = litellm.Router( + model_list=[ + { + "model_name": "copilot-alias", + "litellm_params": {"model": "github_copilot/gpt-4o"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + resolution_attempts: list[str] = [] + + def _oauth_tripwire(model, *args, **kwargs): + resolution_attempts.append(model) + raise AssertionError("get_llm_provider would run the OAuth device flow") + + monkeypatch.setattr(litellm, "get_llm_provider", _oauth_tripwire) + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire) + expected = litellm.get_supported_openai_params(model="gpt-4o", custom_llm_provider="github_copilot") + + with auth_as(): + via_alias = client.get("/utils/supported_openai_params", params={"model": "copilot-alias"}) + via_direct_name = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"}) + + assert via_alias.status_code == 200 + assert via_alias.json() == {"supported_openai_params": expected} + assert via_direct_name.status_code == 200 + assert via_direct_name.json() == {"supported_openai_params": expected} + assert resolution_attempts == [] + + def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch): """Pins ``GET /utils/supported_openai_params`` (error: unknown model).""" def _raise(model): raise Exception("unknown") + monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "get_llm_provider", _raise) with auth_as(): response = client.get("/utils/supported_openai_params", params={"model": "??"}) diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index aa35fd64f18..0fb9b1a6d88 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -19,7 +19,10 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) -from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator +from litellm.proxy.common_utils.model_listing_utils import ( + TeamModelNameTranslator, + configured_display_names, +) from litellm.proxy.proxy_server import ( _get_proxy_model_info, _translate_model_name_for_response, @@ -1391,6 +1394,27 @@ def test_resolve_public_name_respects_legacy_flag(): ) +def test_configured_display_names_keyed_by_response_id(): + """The map is keyed by the public response id while the router lookup uses + the internal routing key, and entries without a configured name are omitted.""" + router = MagicMock() + router.get_configured_display_name = MagicMock( + side_effect=lambda model_name: "Team Sonnet" if model_name == "model_name_team-abc-123_4a6b8" else None + ) + + assert configured_display_names( + entries=[ + ("team-claude-sonnet", "model_name_team-abc-123_4a6b8"), + ("gpt-4o", "gpt-4o"), + ], + llm_router=router, + ) == {"team-claude-sonnet": "Team Sonnet"} + + +def test_configured_display_names_empty_without_router(): + assert configured_display_names(entries=[("gpt-4o", "gpt-4o")], llm_router=None) == {} + + @pytest.mark.asyncio async def test_retrieve_model_by_public_name_returns_200(monkeypatch): """Regression: `GET /v1/models/{public_name}` must NOT 404. The listing diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index c3297ee6ae9..7dd18587df3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -14,6 +14,8 @@ from litellm.proxy.spend_tracking.savings import ( from litellm.router import Router from litellm.types.utils import Usage +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + def _anthropic_costs(model: str) -> tuple[float, float]: info = litellm.get_model_info(model=model, custom_llm_provider="anthropic") diff --git a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py index 99f6f3a9b72..959cb2b1e89 100644 --- a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py +++ b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py @@ -2,6 +2,7 @@ import asyncio import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi.testclient import TestClient @@ -29,6 +30,7 @@ def _make_mock_tts_response(): inner = MagicMock() inner.aiter_bytes = _aiter_bytes inner._hidden_params = {} + inner.response = httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}) async def _resolver(): return inner diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 95067929ac1..b8fb6170d34 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -819,6 +819,94 @@ async def test_should_cap_known_estimate_to_remaining_budget( ) == pytest.approx(0.9) +@pytest.mark.asyncio +async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget( + spend_counter_state, +): + """LIT-5922: with strict enforcement on, a request whose known estimate does + not fit the remaining budget must be rejected before dispatch instead of + having its reservation shrunk to the headroom and admitted, and the counter + must be restored to the pre-request spend.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-known-estimate-fail-closed", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-known-estimate-fail-closed", + value=0.9, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.current_cost == pytest.approx(0.9) + assert exc_info.value.max_budget == pytest.approx(1.0) + assert "Current cost: 0.9, Estimated request cost: 0.6, Max budget: 1.0" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-fail-closed" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits( + spend_counter_state, +): + """0.1 + 0.2 lands a hair above 0.3 in floating point. Strict enforcement + must treat that as fitting the budget, not reject it.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-float-noise", + spend=0.1, + max_budget=0.3, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-fail-closed-float-noise", + value=0.1, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.2, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.2) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-float-noise" + ) == pytest.approx(0.3) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_default_when_output_cap_missing( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 6ea6f208bb5..3e70dee23b7 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2452,6 +2452,96 @@ class TestReadReplicaConnectionParams: assert "DATABASE_URL_READ_REPLICA" not in captured +class TestMaxIdleConnectionLifetimeDefault: + """The proxy defaults `max_idle_connection_lifetime` below common infra idle + timeouts so stale pooled connections are recycled instead of failing requests.""" + + def _config(self, tmp_path, general_settings): + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings})) + return str(config_path) + + def test_default_applied_to_database_and_direct_url(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {}), + direct_url="postgresql://t:t@localhost:5432/t", + ) + + for env_var in ("DATABASE_URL", "DIRECT_URL"): + query = urlparse.parse_qs(urlparse.urlparse(captured[env_var]).query) + assert query["max_idle_connection_lifetime"] == ["60"], env_var + + def test_url_pinned_value_wins_over_default(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {}), + database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300", + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert query["max_idle_connection_lifetime"] == ["300"] + + def test_url_pinned_value_wins_over_config_key(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}), + database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300", + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert query["max_idle_connection_lifetime"] == ["300"] + + def test_config_key_overrides_default(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}), + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert query["max_idle_connection_lifetime"] == ["45"] + + def test_extra_connection_params_override_default(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config( + tmp_path, + {"database_extra_connection_params": {"max_idle_connection_lifetime": 120}}, + ), + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert query["max_idle_connection_lifetime"] == ["120"] + + def test_read_replica_gets_the_default(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {}), + read_replica_url="postgresql://t:t@reader:5432/t", + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) + assert query["max_idle_connection_lifetime"] == ["60"] + + def test_replica_pinned_value_wins(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}), + read_replica_url="postgresql://t:t@reader:5432/t?max_idle_connection_lifetime=200", + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) + assert query["max_idle_connection_lifetime"] == ["200"] + + def test_config_key_reaches_the_read_replica(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}), + read_replica_url="postgresql://t:t@reader:5432/t", + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) + assert query["max_idle_connection_lifetime"] == ["45"] + + def test_idle_lifetime_params_prefers_configured_value(self): + from litellm.proxy.db.db_url_settings import idle_lifetime_params + + assert dict(idle_lifetime_params(45)) == {"max_idle_connection_lifetime": 45} + assert dict(idle_lifetime_params(None)) == {"max_idle_connection_lifetime": 60} + + class TestTokenAuthCliFlags: """`--azure_postgresql_auth` has to reach the URL assembly the same way the env var does.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 43280258153..91fca8f1e27 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10445,6 +10445,75 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): assert before["some_api_key"] != "sk-stored-secret" +@pytest.mark.asyncio +async def test_update_config_field_rejects_out_of_range_alerting_args(monkeypatch): + """Out-of-range alerting_args must be rejected at save time. If they land in the + DB, SlackAlertingArgs raises during the config reload and alerting breaks.""" + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.proxy_server import update_config_general_settings + + monkeypatch.setattr(proxy_server_module, "prisma_client", MagicMock()) + + admin = UserAPIKeyAuth( + api_key="hashed-admin", + user_id="admin-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + with pytest.raises(HTTPException) as exc_info: + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="alerting_args", + field_value={ + "daily_spend_per_user_threshold": -5.0, + "user_spend_check_interval": 20, + }, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.status_code == 400 + error_msg = exc_info.value.detail["error"] + assert "daily_spend_per_user_threshold" in error_msg + assert "user_spend_check_interval" in error_msg + + +@pytest.mark.asyncio +async def test_update_config_field_accepts_valid_alerting_args(monkeypatch): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.proxy_server import update_config_general_settings + + fake = _fake_prisma_with_config({}) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake) + monkeypatch.setattr(litellm, "store_audit_logs", False) + + admin = UserAPIKeyAuth( + api_key="hashed-admin", + user_id="admin-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="alerting_args", + field_value={ + "daily_spend_per_user_threshold": 5.0, + "user_spend_check_interval": 60, + }, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + written = json.loads(fake.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert written["alerting_args"]["daily_spend_per_user_threshold"] == 5.0 + + @pytest.mark.asyncio async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 587be59c550..2b6cfeda2c2 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -111,6 +111,99 @@ def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" +DASHSCOPE_404_BODY = { + "error": { + "message": "The model `does-not-exist` does not exist or you do not have access to it.", + "type": "invalid_request_error", + "param": None, + "code": "model_not_found", + }, + "request_id": "mock-request-id", +} + + +def test_rerank_error_names_provider_and_keeps_body(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the rerank error path mapping with the unresolved provider param: + a provider 404 surfaced as 'None - ' instead of naming the provider and its error body.""" + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + mock_route = respx_mock.post("https://dashscope.example/v1/reranks") + mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY) + + with pytest.raises(litellm.NotFoundError) as exc_info: + litellm.rerank( + model="dashscope/does-not-exist", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + api_base="https://dashscope.example/v1", + ) + + assert mock_route.called + assert "DashscopeException" in str(exc_info.value) + assert "does not exist or you do not have access to it" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.MockRouter, monkeypatch): + """Regression for arerank's bare re-raise: provider errors escaped as raw + provider exception classes instead of the mapped litellm exception contract.""" + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + mock_route = respx_mock.post("https://dashscope.example/v1/reranks") + mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY) + + with pytest.raises(litellm.NotFoundError) as exc_info: + await litellm.arerank( + model="dashscope/does-not-exist", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + api_base="https://dashscope.example/v1", + ) + + assert mock_route.called + assert "DashscopeException" in str(exc_info.value) + assert "does not exist or you do not have access to it" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_arerank_declared_authenticating_provider_skips_resolution(monkeypatch): + """Regression for the event-loop hazard in arerank's provider pre-resolution: + get_llm_provider runs the blocking OAuth device flow for github_copilot/chatgpt, + so arerank must adopt the declared provider instead of resolving it, while the + except path still maps with that declared provider.""" + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + resolution_calls = [] + + def record_resolution(*args, **kwargs): + resolution_calls.append((args, kwargs)) + return "gpt-4o", "github_copilot", None, None + + def rerank_raises_provider_error(*args, **kwargs): + raise BaseLLMException(status_code=401, message='{"error":"bad key"}') + + monkeypatch.setattr(litellm, "get_llm_provider", record_resolution) + monkeypatch.setattr("litellm.rerank_api.main.rerank", rerank_raises_provider_error) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + await litellm.arerank( + model="github_copilot/gpt-4o", + query=MARKER_QUERY, + documents=[MARKER_DOC], + ) + + assert resolution_calls == [] + assert "Github_copilotException" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + @pytest.mark.asyncio async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index bb59e576568..b2b8eb5da80 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -648,6 +648,72 @@ class TestLiteLLMCompletionResponsesConfig: assert responses_api_response.status == "incomplete" + def test_tool_call_only_response_emits_no_null_text_message_item(self): + """A tool-calls-only turn (message content None, e.g. from Anthropic) + must not emit a message output item whose output_text has text null. + OpenAI rejects such an item on replay with + "Invalid type for 'input[..].content[..].text': expected a string, but + got null instead." Native OpenAI tool-only turns carry no message item.""" + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="toolu_01OnlyToolCall", + type="function", + function=Function(name="get_weather", arguments='{"city": "SF"}'), + ) + ], + ), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="what's the weather in SF?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) + + output_types = [item.type for item in responses_api_response.output] + assert "message" not in output_types + assert "function_call" in output_types + + def test_content_bearing_response_still_emits_message_item(self): + """Turns with real text content must keep their message output item.""" + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="claude-sonnet-4-5", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="It is sunny.", role="assistant"), + ) + ], + ) + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="what's the weather in SF?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) + + message_items = [item for item in responses_api_response.output if item.type == "message"] + assert len(message_items) == 1 + assert message_items[0].content[0].text == "It is sunny." + def test_transform_chat_completion_response_preserves_hidden_params(self): """Test that _hidden_params from chat completion response are preserved in responses API response""" # Setup @@ -985,6 +1051,55 @@ class TestFunctionCallTransformation: assert result[0]["tool_calls"][0]["function"]["arguments"] == "{}" + def test_function_call_transformation_json_encodes_object_arguments(self): + """A decoded arguments object must be JSON-encoded, not str()'d. + + Clients and providers sometimes send `arguments` as an object rather + than a JSON string; `str()` on a dict produces a Python repr with + single quotes, which downstream JSON parsers reject with errors like + "Expecting ',' delimiter". + """ + function_call_item = { + "type": "function_call", + "name": "shell", + "arguments": {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]}, + "call_id": "call_123", + "id": "call_123", + "status": "completed", + } + + result = LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( + function_call=function_call_item + ) + + arguments = result[0].get("tool_calls", [])[0].get("function", {}).get("arguments") + assert json.loads(arguments) == {"command": "ls", "timeout": 30, "flags": ["-l", "-a"]} + assert "'" not in arguments + + def test_create_tool_call_chunk_json_encodes_object_arguments(self): + """Cached tool_call definitions with object arguments stay valid JSON.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={ + "id": "call_456", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls"}}, + }, + tool_call_id="call_456", + index=0, + ) + + assert json.loads(chunk["function"]["arguments"]) == {"command": "ls"} + + def test_create_tool_call_chunk_keeps_empty_arguments_default(self): + """Missing arguments still fall back to an empty JSON object.""" + chunk = LiteLLMCompletionResponsesConfig._create_tool_call_chunk( + tool_use_definition={"id": "call_789", "type": "function", "function": {"name": "shell"}}, + tool_call_id="call_789", + index=0, + ) + + assert chunk["function"]["arguments"] == "{}" + def test_complete_input_transformation_with_function_calls(self): """Test the complete transformation with the exact input from the issue""" test_input = [ @@ -3294,6 +3409,7 @@ class TestEnsureOutputItemContentPartAdded: iterator._pending_tool_events = [] iterator._tool_output_index_by_call_id = {} iterator._tool_args_by_call_id = {} + iterator._tool_item_id_by_call_id = {} iterator._tool_call_id_by_index = {} iterator._ambiguous_tool_call_indexes = set() iterator._next_tool_output_index = 1 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 823f656ddc5..4a03913f55a 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -10,6 +10,7 @@ before response.completed, and that every event of a bridged stream carries the spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ +import json from unittest.mock import AsyncMock, MagicMock import pytest @@ -131,7 +132,7 @@ def test_tool_call_delta_is_emitted_as_responses_events(): evt2 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) assert evt2 is not None assert evt2.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - assert evt2.item_id == "call_1" + assert evt2.item_id == "fc_call_1" assert evt2.output_index == 1 # The delta will be a chunk of the arguments, not the full arguments assert len(evt2.delta) <= 10 # Chunks are max 10 characters @@ -196,7 +197,7 @@ def test_tool_calls_present_only_in_final_response_are_emitted_before_completed( # The last event should be FUNCTION_CALL_ARGUMENTS_DONE assert evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE - assert evt.item_id == "call_2" + assert evt.item_id == "fc_call_2" assert evt.output_index == 1 assert evt.arguments == '{"y":2}' @@ -290,7 +291,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): # Verify each delta is at most 10 characters for evt in delta_events: assert len(evt.delta) <= 10 - assert evt.item_id == "call_test" + assert evt.item_id == "fc_call_test" assert evt.output_index == 1 assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__ @@ -348,7 +349,8 @@ def test_tool_call_delta_without_id_uses_index_mapping(): if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED ] assert len(output_item_added_events) == 1 - assert output_item_added_events[0].item.id == "call_abc123" + assert output_item_added_events[0].item.id == "fc_call_abc123" + assert output_item_added_events[0].item.call_id == "call_abc123" def test_parallel_tool_calls_without_ids_use_index_mapping(): @@ -403,8 +405,8 @@ def test_parallel_tool_calls_without_ids_use_index_mapping(): arguments_by_call_id.setdefault(evt.item_id, "") arguments_by_call_id[evt.item_id] += evt.delta - assert arguments_by_call_id["call_a"] == '{"x":1}' - assert arguments_by_call_id["call_b"] == '{"y":2}' + assert arguments_by_call_id["fc_call_a"] == '{"x":1}' + assert arguments_by_call_id["fc_call_b"] == '{"y":2}' def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): @@ -460,10 +462,10 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): arguments_by_call_id.setdefault(evt.item_id, "") arguments_by_call_id[evt.item_id] += evt.delta - assert arguments_by_call_id["call_a"] == '{"a":' - assert arguments_by_call_id["call_b"] == '{"b":' - assert arguments_by_call_id["call_a"] != '{"a":1}' - assert arguments_by_call_id["call_b"] != '{"b":1}' + assert arguments_by_call_id["fc_call_a"] == '{"a":' + assert arguments_by_call_id["fc_call_b"] == '{"b":' + assert arguments_by_call_id["fc_call_a"] != '{"a":1}' + assert arguments_by_call_id["fc_call_b"] != '{"b":1}' @pytest.mark.asyncio @@ -523,3 +525,91 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids assert len(set(response_ids)) == 1 assert response_ids[0].startswith("resp_") + + +def test_object_tool_call_arguments_stream_as_valid_json(): + """A provider that sends decoded object arguments must still stream valid JSON. + + `str()` on a dict yields a Python repr with single quotes, which clients + parsing function_call_arguments reject with errors like + "Expecting ',' delimiter". + """ + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + iterator._queue_tool_call_delta_events( + [ + { + "index": 0, + "id": "call_obj", + "type": "function", + "function": {"name": "shell", "arguments": {"command": "ls", "flags": ["-l"]}}, + } + ] + ) + + streamed_arguments = "".join( + evt.delta + for evt in iterator._pending_tool_events + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ) + + assert json.loads(streamed_arguments) == {"command": "ls", "flags": ["-l"]} + + +def test_streamed_anthropic_tool_call_events_correlate_on_normalized_item_id(): + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + + response = ModelResponse( + id="resp-anthropic", + created=123, + model="test-model", + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "toolu_01AbCdEf", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ], + }, + } + ], + ) + iterator.litellm_model_response = response + + events = [] + while True: + evt = iterator.common_done_event_logic(sync_mode=True) + events.append(evt) + if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + break + + added = [e for e in events if e.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] + deltas = [e for e in events if e.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] + dones = [e for e in events if e.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE] + item_dones = [e for e in events if e.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] + + assert len(added) == 1 and len(dones) == 1 and len(item_dones) == 1 and deltas + assert added[0].item.id == "fc_toolu_01AbCdEf" + assert added[0].item.call_id == "toolu_01AbCdEf" + assert item_dones[0].item.id == "fc_toolu_01AbCdEf" + assert item_dones[0].item.call_id == "toolu_01AbCdEf" + for evt in deltas + dones: + assert evt.item_id == added[0].item.id diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/test_litellm/responses/test_custom_tool_call.py index c605ef24934..5122c1c1d67 100644 --- a/tests/test_litellm/responses/test_custom_tool_call.py +++ b/tests/test_litellm/responses/test_custom_tool_call.py @@ -20,6 +20,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( from litellm.responses.litellm_completion_transformation.custom_tools import ( extract_custom_tool_names, is_custom_tool_call, + openai_shaped_tool_call_item_id, unwrap_custom_tool_arguments, build_tool_call_item_kwargs, convert_custom_tool_to_function_tool, @@ -129,6 +130,41 @@ class TestCustomToolUtilities: assert kwargs["arguments"] == raw assert "input" not in kwargs + def test_openai_shaped_tool_call_item_id_prefixes_foreign_ids(self): + """Anthropic-style tool ids must be normalized to OpenAI's item id + shapes (fc/ctc prefixes) so replaying the item to OpenAI does not 400 + with "Expected an ID that begins with 'fc'".""" + assert openai_shaped_tool_call_item_id("function_call", "toolu_01Abc") == "fc_toolu_01Abc" + assert openai_shaped_tool_call_item_id("function_call", "srvtoolu_01Xyz") == "fc_srvtoolu_01Xyz" + assert openai_shaped_tool_call_item_id("custom_tool_call", "toolu_01Abc") == "ctc_toolu_01Abc" + assert openai_shaped_tool_call_item_id("function_call", "fc_already") == "fc_already" + assert openai_shaped_tool_call_item_id("custom_tool_call", "ctc_already") == "ctc_already" + assert openai_shaped_tool_call_item_id("function_call", "") == "" + assert openai_shaped_tool_call_item_id("message", "toolu_01Abc") == "toolu_01Abc" + + def test_build_tool_call_item_kwargs_normalizes_item_id_keeps_call_id(self): + """The streaming item id gets the OpenAI shape while call_id stays raw + so tool_result pairing (which keys off call_id) keeps working.""" + function_kwargs = build_tool_call_item_kwargs( + call_id="toolu_01Abc", + name="get_weather", + arguments_or_input="{}", + status="completed", + custom_tool_names=set(), + ) + assert function_kwargs["id"] == "fc_toolu_01Abc" + assert function_kwargs["call_id"] == "toolu_01Abc" + + custom_kwargs = build_tool_call_item_kwargs( + call_id="toolu_01Def", + name="apply_patch", + arguments_or_input=json.dumps({"content": "patch"}), + status="completed", + custom_tool_names={"apply_patch"}, + ) + assert custom_kwargs["id"] == "ctc_toolu_01Def" + assert custom_kwargs["call_id"] == "toolu_01Def" + def test_unwrap_custom_tool_arguments_oversized_returns_raw(self): """Arguments larger than the safety cap are returned unchanged to avoid OOM on JSON parsing a pathologically large string.""" @@ -293,6 +329,52 @@ class TestTransformationCustomTools: assert item.name == "regular_tool" assert item.arguments == json.dumps({"param": "value"}) + def test_transform_anthropic_tool_call_ids_get_openai_item_id_shape(self): + """Anthropic tool ids (toolu_/srvtoolu_) surfacing through the bridge + must be emitted with fc/ctc-prefixed item ids so a Responses client can + replay them to OpenAI verbatim, while call_id stays raw for pairing.""" + from litellm.types.utils import ModelResponse, Choices, Message, ChatCompletionMessageToolCall, Function + + client_call = ChatCompletionMessageToolCall( + id="toolu_01ClientCall", + type="function", + function=Function(name="get_weather", arguments=json.dumps({"city": "SF"})), + ) + server_call = ChatCompletionMessageToolCall( + id="srvtoolu_01ServerCall", + type="function", + function=Function(name="web_search", arguments=json.dumps({"query": "zig"})), + ) + custom_call = ChatCompletionMessageToolCall( + id="toolu_01CustomCall", + type="function", + function=Function(name="apply_patch", arguments=json.dumps({"content": "patch content"})), + ) + + message = Message(role="assistant", content=None, tool_calls=[client_call, server_call, custom_call]) + choices = [Choices(index=0, message=message, finish_reason="tool_calls")] + response = ModelResponse( + id="test_response", choices=choices, created=1234567890, model="claude-sonnet-4-5", object="chat.completion" + ) + responses_api_request = { + "tools": [{"type": "custom", "name": "apply_patch"}, {"type": "function", "name": "get_weather"}] + } + + result = LiteLLMCompletionResponsesConfig.transform_chat_completion_tools_to_responses_tools( + response, responses_api_request=responses_api_request + ) + + assert [item.id for item in result] == [ + "fc_toolu_01ClientCall", + "fc_srvtoolu_01ServerCall", + "ctc_toolu_01CustomCall", + ] + assert [item.call_id for item in result] == [ + "toolu_01ClientCall", + "srvtoolu_01ServerCall", + "toolu_01CustomCall", + ] + def test_transform_mixed_tool_calls(self): """Test transformation with both custom and regular tool calls.""" from litellm.types.utils import ModelResponse, Choices, Message, ChatCompletionMessageToolCall, Function diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 677faf7f655..9edcaaef034 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -326,3 +326,55 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0 assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params + + +def _responses_api_response_with_usage() -> ResponsesAPIResponse: + from litellm.types.llms.openai import ResponseAPIUsage + + return ResponsesAPIResponse( + id="resp_lit6427", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="mantle-claude", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=20, output_tokens=60, total_tokens=80), + ) + + +def test_stamp_responses_usage_cost_stamps_computed_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + +def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + setattr(response.usage, "cost", 0.5) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_survives_calculator_failure(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.side_effect = RuntimeError("cost map unavailable") + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) is None diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index b3a2bdda53c..60433921de6 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -598,6 +598,45 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh assert filtered == healthy_deployments +@pytest.mark.asyncio +async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): + user_key = "user-key-order-fallback" + stable_model_map_key = "claude-sonnet-4-5@20250929" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model_id": "deployment-1"}) + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + healthy_deployments = [ + { + "model_name": stable_model_map_key, + "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": stable_model_map_key, + "litellm_params": { + "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" + }, + "model_info": {"id": "deployment-2"}, + }, + ] + + filtered = await callback.async_filter_deployments( + model="some-router-model-group", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"_target_order": 2, "metadata": {"user_api_key_hash": user_key}}, + parent_otel_span=None, + ) + + assert filtered == healthy_deployments + cache.async_get_cache.assert_not_called() + + @pytest.mark.asyncio async def test_async_user_key_affinity_ttl_expiry_allows_reroute(): """ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index f54a1cfa284..79ae00e155c 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -150,6 +150,25 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): assert filtered == [deployments[1]] +@pytest.mark.asyncio +async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + messages = _messages(word_count=5000) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"_target_order": 2}, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 8336926c050..894b2d9e74f 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -11,7 +11,10 @@ from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _trigger_cooldown_for_failed_deployment, fallback_attempt_key, + clear_pre_routing_selection, get_fallback_model_group, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) @@ -1090,3 +1093,119 @@ async def test_run_async_fallback_preserves_original_model_group_on_nested_fallb metadata = router.received_kwargs["metadata"] assert metadata["attempted_fallbacks"] == 2 assert metadata["original_model_group"] == "primary-model" + + +class TestPreRoutingSelectionCarriesToFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup kept using the router name, so the tier's configured chain never ran.""" + + def test_selection_is_recorded_in_the_metadata_bucket(self): + kwargs = {"model": "smart-router", "metadata": {}} + record_pre_routing_selection(kwargs, "tier1") + assert kwargs["metadata"]["pre_routing_selected_model"] == "tier1" + assert get_pre_routing_selection(kwargs) == "tier1" + + def test_selection_is_recorded_in_the_litellm_metadata_bucket(self): + kwargs = {"model": "smart-router", "litellm_metadata": {}} + record_pre_routing_selection(kwargs, "tier2") + assert get_pre_routing_selection(kwargs) == "tier2" + + def test_a_bucket_survives_the_kwargs_copy_that_fallbacks_run_on(self): + """The bucket is shared by reference, which is the whole reason this works.""" + outer = {"model": "smart-router", "metadata": {}} + inner = {**outer} + record_pre_routing_selection(inner, "tier1") + assert get_pre_routing_selection(outer) == "tier1" + + def test_no_selection_reads_as_none(self): + assert get_pre_routing_selection({"model": "smart-router", "metadata": {}}) is None + assert get_pre_routing_selection({"model": "smart-router"}) is None + + def test_missing_kwargs_is_a_no_op(self): + """A caller with no kwargs must not raise, and must not leak the selection anywhere.""" + record_pre_routing_selection(None, "tier1") + + assert get_pre_routing_selection({}) is None + + def test_a_non_dict_bucket_is_ignored(self): + kwargs = {"model": "smart-router", "metadata": "not-a-dict"} + record_pre_routing_selection(kwargs, "tier1") + assert get_pre_routing_selection(kwargs) is None + + def test_fallbacks_resolve_against_the_selected_tier(self): + """The lookup the router performs, keyed on the tier rather than the router name.""" + fallbacks = [{"tier1": ["backup-a", "backup-b"]}, {"tier2": ["backup-c"]}] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="tier1")[0] == ["backup-a", "backup-b"] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="smart-router")[0] is None + + +class TestPreRoutingSelectionIsPerHop: + """#38832 review: the buckets also carry whatever the caller sent, and a fallback hop + inherits the previous hop's tier, so a hop must start without a selection.""" + + def test_a_caller_supplied_selection_is_dropped(self): + kwargs = {"model": "plain", "metadata": {"pre_routing_selected_model": "tier1"}} + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + assert "pre_routing_selected_model" not in kwargs["metadata"] + + def test_both_buckets_are_cleared(self): + kwargs = { + "metadata": {"pre_routing_selected_model": "tier1"}, + "litellm_metadata": {"pre_routing_selected_model": "tier2"}, + } + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + + def test_the_rest_of_the_bucket_is_left_alone(self): + kwargs = {"metadata": {"pre_routing_selected_model": "tier1", "tags": ["a"]}} + + clear_pre_routing_selection(kwargs) + + assert kwargs["metadata"] == {"tags": ["a"]} + + def test_clearing_is_a_no_op_without_a_usable_bucket(self): + kwargs = {"model": "plain", "metadata": "not-a-dict"} + + clear_pre_routing_selection(None) + clear_pre_routing_selection(kwargs) + + assert kwargs == {"model": "plain", "metadata": "not-a-dict"} + + def test_a_selection_recorded_after_clearing_is_kept(self): + """Clearing runs before routing, so the hook's own write must survive it.""" + kwargs = {"model": "smart-router", "metadata": {"pre_routing_selected_model": "stale"}} + + clear_pre_routing_selection(kwargs) + record_pre_routing_selection(kwargs, "tier1") + + assert get_pre_routing_selection(kwargs) == "tier1" + + +class TestOrderedFallbackLookupGroups: + def test_tier_first_then_requested_group_deduped(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier1"}} + assert fallback_lookup_groups(kwargs, "smart-router") == ("tier1", "smart-router") + assert fallback_lookup_groups(kwargs, "tier1") == ("tier1",) + assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) + assert fallback_lookup_groups({}, None) == () + + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): + from litellm.router_utils.fallback_event_handlers import ( + get_fallback_model_group_for_lookup_groups, + ) + + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}, {"*": ["backup-c"]}] + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier1", "smart-router")) == (["backup-a"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "smart-router")) == (["backup-b"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "no-such")) == (["backup-c"], 2) + assert get_fallback_model_group_for_lookup_groups([{"tier1": ["backup-a"]}], ("no", "nope")) == (None, None) diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/test_litellm/router_utils/test_pattern_match_deployments.py new file mode 100644 index 00000000000..795d448ef5f --- /dev/null +++ b/tests/test_litellm/router_utils/test_pattern_match_deployments.py @@ -0,0 +1,78 @@ +"""Behavior pins for ``litellm/router_utils/pattern_match_deployments.py``.""" + +from __future__ import annotations + +from litellm.router_utils import pattern_match_deployments +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + +def _wildcard_deployment(model_name: str) -> dict: + return {"model_name": model_name, "litellm_params": {"model": model_name}} + + +def _matched_models(matches: list[dict] | None) -> list[str]: + return [deployment["litellm_params"]["model"] for deployment in matches or []] + + +def test_get_pattern_never_resolves_declared_authenticating_providers(monkeypatch): + """Regression: resolving a github_copilot/chatgpt name through ``get_llm_provider`` runs the + provider's OAuth device flow; the auth layer walks every wildcard router on every request, so + a single metadata lookup for an unserved name would block the proxy's event loop.""" + resolution_attempts: list[str] = [] + + def _oauth_tripwire(model, *args, **kwargs): + resolution_attempts.append(model) + raise AssertionError("get_llm_provider would run the OAuth device flow") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire) + + unmatched_router = PatternMatchRouter() + unmatched_router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*")) + assert unmatched_router.get_pattern("github_copilot/gpt-4o") is None + + matched_router = PatternMatchRouter() + matched_router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*")) + assert _matched_models(matched_router.get_pattern("github_copilot/gpt-4o")) == ["github_copilot/gpt-4o"] + assert _matched_models(matched_router.get_pattern("gpt-4o", custom_llm_provider="github_copilot")) == [ + "github_copilot/gpt-4o" + ] + + assert resolution_attempts == [] + + +def test_get_pattern_bare_provider_name_never_matches_that_providers_wildcard(monkeypatch): + """Regression: a bare ``github_copilot`` adopted itself as its provider and retried as + ``github_copilot/github_copilot``, false-matching the wildcard for a name no deployment serves.""" + + def _unknown_provider(model, *args, **kwargs): + raise ValueError(f"unknown provider for {model}") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider) + router = PatternMatchRouter() + router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*")) + assert router.get_pattern("github_copilot") is None + + +def test_get_pattern_missing_model_returns_none(monkeypatch): + """Regression: a request without a model reaches the auth layer's pattern walk as ``None``; the + declared-provider guard raised ``TypeError`` where the old inline resolve swallowed every + resolver error, so the proxy's missing-model 400 became a crash.""" + + def _unknown_provider(model, *args, **kwargs): + raise ValueError(f"unknown provider for {model}") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider) + router = PatternMatchRouter() + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + assert router.get_pattern(None) is None + + +def test_get_pattern_still_resolves_unqualified_names(monkeypatch): + monkeypatch.setattr( + pattern_match_deployments, + "get_llm_provider", + lambda model, **kwargs: (model, "openai", None, None), + ) + router = PatternMatchRouter() + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"] diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 6b3312b5cc4..f7d95ecda01 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -26,9 +26,7 @@ import pytest @pytest.fixture(scope="module") def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) + json_path = os.path.join(os.path.dirname(__file__), "../../model_prices_and_context_window.json") with open(json_path) as f: return json.load(f) @@ -51,21 +49,14 @@ def test_usgov_sonnet_4_5_pricing(model_data, model_key): info = model_data[model_key] assert info["input_cost_per_token"] == 3.6e-06, ( - f"{model_key}: input_cost_per_token should be $3.60/MTok " - f"(got {info['input_cost_per_token']})" + f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})" ) - assert ( - info["output_cost_per_token"] == 1.8e-05 - ), f"{model_key}: output_cost_per_token should be $18.00/MTok" - assert ( - info["cache_creation_input_token_cost"] == 4.5e-06 - ), f"{model_key}: 5m cache write should be $4.50/MTok" - assert ( - info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06 - ), f"{model_key}: 1h cache write should be $7.20/MTok" - assert ( - info["cache_read_input_token_cost"] == 3.6e-07 - ), f"{model_key}: cache read should be $0.36/MTok" + assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok" + assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok" + assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, ( + f"{model_key}: 1h cache write should be $7.20/MTok" + ) + assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok" def test_usgov_carries_20_percent_premium_over_global(model_data): @@ -84,9 +75,7 @@ def test_usgov_carries_20_percent_premium_over_global(model_data): "cache_read_input_token_cost", ): ratio = usgov_info[field] / global_info[field] - assert ( - abs(ratio - 1.2) < 1e-9 - ), f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" # The us-gov.anthropic.* cross-region inference profile is the only us-gov @@ -112,9 +101,7 @@ def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, ex """ info = model_data[USGOV_CROSS_REGION_KEY] assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}" - assert ( - info[field] == expected - ), f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" + assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" def test_usgov_cross_region_above_200k_ratio_to_global(model_data): @@ -127,6 +114,176 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data): usgov_info = model_data[USGOV_CROSS_REGION_KEY] for field in EXPECTED_USGOV_ABOVE_200K: ratio = usgov_info[field] / global_info[field] - assert ( - abs(ratio - 1.2) < 1e-9 - ), f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + + +CLAUDE_GOV_EXPECTED = { + "anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + }, + "anthropic.claude-opus-4-8": { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 3e-05, + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + }, +} + + +USGOV_CLAUDE_KEY_TEMPLATES = { + "bedrock/us-gov-east-1/{base_key}": "bedrock", + "bedrock/us-gov-west-1/{base_key}": "bedrock", + "us-gov.{base_key}": "bedrock_converse", +} + + +@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) +@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) +def test_usgov_claude_sonnet5_opus48_pricing(model_data, key_template, expected_provider, base_key): + """Sonnet 5 and Opus 4.8 gov entries, both in-region keys and the us-gov. + geo inference profile the model cards list for GovCloud, must match the + rates AWS publishes on the Bedrock pricing page (1.2x global). + """ + gov_key = key_template.format(base_key=base_key) + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + assert info["litellm_provider"] == expected_provider + for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + ratio = info[field] / model_data[base_key][field] + assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2" + + +CONVERSE_GOV_EXPECTED = { + "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07), + "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07), + "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07), + "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07), + "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07), +} + + +@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED) +@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) +def test_usgov_converse_model_pricing(model_data, region, base_key): + """Nemotron and gpt-oss gov entries must match the AWS Bedrock offer file, + which prices both GovCloud regions identically at 1.2x commercial. + """ + gov_key = f"bedrock/{region}/{base_key}" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key] + assert info["input_cost_per_token"] == expected_input + assert info["output_cost_per_token"] == expected_output + assert info["litellm_provider"] == "bedrock" + base = model_data[base_key] + assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9 + assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9 + + +def test_usgov_west_llama3_8b_output_price_fixed(model_data): + """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok); + the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model + in us-gov-west-1 only, so there is no east entry to check. + """ + info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"] + assert info["input_cost_per_token"] == 3e-07 + assert info["output_cost_per_token"] == 6e-07 + + +MANTLE_GOV_TIERED_EXPECTED = { + "openai.gpt-5.6-luna": { + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, + }, + "openai.gpt-5.6-terra": { + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, + }, +} + + +@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED) +def test_usgov_west_mantle_terra_luna_pricing(model_data, model): + """Terra and Luna carry 1.2x commercial across every tier in the + us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them. + """ + gov_key = f"bedrock_mantle/us-gov-west-1/{model}" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + assert info["litellm_provider"] == "bedrock_mantle" + assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data + + +@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) +def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region): + """gpt-5.4 gov rates come from the offer file, which publishes only the + standard tier in GovCloud: no long-context SKUs exist there, unlike commercial. + """ + gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + assert info["input_cost_per_token"] == 3.3e-06 + assert info["cache_read_input_token_cost"] == 3.3e-07 + assert info["output_cost_per_token"] == 1.98e-05 + assert not any(field.endswith("_above_272k_tokens") for field in info) + + +def test_usgov_mantle_grok_4_3_west_only(model_data): + """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer + file carries grok-4.6 instead. + """ + info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + assert info["input_cost_per_token"] == 1.5e-06 + assert info["output_cost_per_token"] == 3e-06 + assert info["cache_read_input_token_cost"] == 2.4e-07 + assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data + + +AZURE_GOV_EXPECTED = { + "azure/us-gov/gpt-5.1": { + "input_cost_per_token": 1.71875e-06, + "cache_read_input_token_cost": 1.71875e-07, + "output_cost_per_token": 1.375e-05, + }, + "azure/us-gov/o3-mini": { + "input_cost_per_token": 1.513e-06, + "cache_read_input_token_cost": 7.57e-07, + "output_cost_per_token": 6.05e-06, + }, + "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07}, + "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08}, +} + + +@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED) +def test_azure_usgov_pricing(model_data, gov_key): + """Azure Government meters from the Azure retail prices API + (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government + retirement schedule is published, so these entries carry no deprecation_date. + """ + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + for field, expected in AZURE_GOV_EXPECTED[gov_key].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + assert info["litellm_provider"] == "azure" + assert "deprecation_date" not in info diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py index 99c59ffa58e..3ecf94602d9 100644 --- a/tests/test_litellm/test_claude_fable_5_config.py +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -1,5 +1,5 @@ """ -Validate Claude Fable 5 model configuration entries. +Validate Claude Fable 5 and Claude Fable 5.1 model configuration entries. Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only API surface as Opus 4.7/4.8. The cost-map entries below are what make the model @@ -210,6 +210,156 @@ def test_adaptive_thinking_detected_for_fable_5(local_model_cost_map, model): assert AnthropicModelInfo._is_adaptive_thinking_model(model, "anthropic") is True +FABLE_5_1_VARIANTS = ( + "claude-fable-5-1", + "anthropic.claude-fable-5-1", + "global.anthropic.claude-fable-5-1", + "us.anthropic.claude-fable-5-1", + "eu.anthropic.claude-fable-5-1", + "vertex_ai/claude-fable-5-1", + "vertex_ai/claude-fable-5-1@default", + "azure_ai/claude-fable-5-1", +) + + +def test_fable_5_1_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = [ + ("claude-fable-5-1", "anthropic"), + ("anthropic.claude-fable-5-1", "bedrock_converse"), + ("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"), + ("azure_ai/claude-fable-5-1", "azure_ai"), + ] + + for model_name, provider in expected_models: + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + assert info["input_cost_per_token"] == 1e-05 + assert info["output_cost_per_token"] == 5e-05 + assert info["cache_creation_input_token_cost"] == 1.25e-05 + assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 + + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_assistant_prefill"] is False + assert info["supports_forced_tool_use"] is False + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + assert info["prompt_cache_min_tokens"] == 512 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map): + """Fable 5.1 prices cache hits at 0.025x base input instead of the usual + 0.1x, so copying Fable 5's cache-read price overcharges every cache hit 4x.""" + for model_name in FABLE_5_1_VARIANTS: + info = cost_map[model_name] + geo_premium = model_name.startswith(("us.", "eu.")) + expected = 2.75e-07 if geo_premium else 2.5e-07 + assert info["cache_read_input_token_cost"] == expected, model_name + assert info["cache_read_input_token_cost"] == pytest.approx( + info["input_cost_per_token"] * 0.025 + ), model_name + + +def test_fable_5_1_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + expected_models = { + "global.anthropic.claude-fable-5-1": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 2.5e-07, + }, + "us.anthropic.claude-fable-5-1": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 2.75e-07, + }, + "eu.anthropic.claude-fable-5-1": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 2.75e-07, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["bedrock_output_config_effort_ceiling"] == "xhigh" + for key, value in expected.items(): + assert info[key] == value + + +def test_fable_5_1_geo_multiplier_without_fast_mode(): + """Fable 5.1 has no fast mode, so a ``fast`` key here would misprice + ``speed='fast'`` requests.""" + model_data = _load_root_cost_map() + assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1} + + +def test_fable_5_1_present_in_bundled_backup(): + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in FABLE_5_1_VARIANTS: + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +def test_fable_5_1_registered_for_bedrock_converse(): + assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS + + +def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map): + info = litellm.get_model_info(model="claude-fable-5-1") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5-1", + "anthropic/claude-fable-5-1", + "anthropic.claude-fable-5-1", + "bedrock/us.anthropic.claude-fable-5-1", + "bedrock/invoke/eu.anthropic.claude-fable-5-1", + "bedrock/global.anthropic.claude-fable-5-1", + "vertex_ai/claude-fable-5-1", + "azure_ai/claude-fable-5-1", + ], +) +def test_adaptive_thinking_detected_for_fable_5_1(local_model_cost_map, model): + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo._is_adaptive_thinking_model(model, "anthropic") is True + + @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py index 9ca4515239a..e33bcfb8378 100644 --- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -75,6 +75,22 @@ def test_additional_current_models_are_present(): assert entry["output_cost_per_token"] > 0 +@pytest.mark.parametrize( + "key, published_price_per_audio_minute", + [ + ("cloudflare/@cf/openai/whisper", 0.00045), + ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051), + ], +) +def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute): + entry = litellm.model_cost[key] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "audio_transcription" + assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] + assert entry["output_cost_per_second"] == 0.0 + assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60) + + def test_root_and_backup_have_identical_cloudflare_keys(): if not os.path.exists(ROOT_MAP): pytest.skip("root cost map only ships in source checkouts") diff --git a/tests/test_litellm/test_dockerfile_apk_repository.py b/tests/test_litellm/test_dockerfile_apk_repository.py new file mode 100644 index 00000000000..cbd772defbf --- /dev/null +++ b/tests/test_litellm/test_dockerfile_apk_repository.py @@ -0,0 +1,52 @@ +""" +Static checks on the root Dockerfile's apk repository configuration. + +The base image (cgr.dev/chainguard/wolfi-base) only configures the +authenticated Chainguard apk repo (https://apk.cgr.dev/chainguard) in +/etc/apk/repositories, which requires a Chainguard enterprise subscription. +Anyone pulling the published litellm image and running `apk add` inside it +hits SSL/auth failures with no fallback repo configured, so nothing can be +installed. See https://github.com/BerriAI/litellm/issues/33518 +""" + +import os +import re + +import pytest + +DOCKERFILE_PATH = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "Dockerfile", +) + + +def _runtime_stage(dockerfile_text: str) -> str: + """Return the contents of the final `FROM ... AS runtime` build stage.""" + match = re.search(r"^FROM .*\bAS runtime\b(.*)\Z", dockerfile_text, re.MULTILINE | re.DOTALL) + assert match, "Dockerfile has no `FROM ... AS runtime` stage" + return match.group(1) + + +@pytest.mark.skipif( + not os.path.exists(DOCKERFILE_PATH), + reason="Dockerfile not present in this checkout", +) +def test_runtime_stage_adds_public_wolfi_repo(): + """The runtime stage must add the public Wolfi apk repo so `apk add` + works for users without a Chainguard enterprise subscription.""" + with open(DOCKERFILE_PATH, "r", encoding="utf-8") as f: + contents = f.read() + + runtime_stage = _runtime_stage(contents) + + assert re.search( + r"echo\s+[\"']?https://packages\.wolfi\.dev/os[\"']?\s*>>\s*/etc/apk/repositories", + runtime_stage, + ), ( + "Runtime stage must append the public Wolfi apk repo " + '(RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories) ' + "so `apk add` works without Chainguard enterprise credentials. " + "See https://github.com/BerriAI/litellm/issues/33518" + ) diff --git a/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py b/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py new file mode 100644 index 00000000000..44572aed08e --- /dev/null +++ b/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py @@ -0,0 +1,56 @@ +""" +Static checks that every proxy Docker image installs the `bedrock-realtime` extra. + +Bedrock Nova Sonic speech-to-speech (`/v1/realtime`) needs `aws-sdk-bedrock-runtime`, +which only ships in the `bedrock-realtime` extra. An image whose `uv sync` stages +omit the extra fails every Nova Sonic realtime session with +"Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime". +""" + +import os +import re +from typing import Final + +import pytest + +REPO_ROOT: Final = os.path.join(os.path.dirname(__file__), "..", "..") + +PROXY_DOCKERFILES: Final = ( + "Dockerfile", + os.path.join("docker", "Dockerfile.non_root"), + os.path.join("docker", "Dockerfile.database"), + os.path.join("gateway", "Dockerfile"), +) + +CONTINUED_LINE_RE: Final = re.compile(r"(?:\\\n|[^\n])+") +UV_SYNC_BOUNDARY_RE: Final = re.compile(r"(?=uv sync)") + + +def _uv_sync_invocations(dockerfile_text: str) -> tuple[str, ...]: + """Return each `uv sync ...` command, split apart when one RUN holds several (if/else branches).""" + return tuple( + part + for line in CONTINUED_LINE_RE.finditer(dockerfile_text) + for part in UV_SYNC_BOUNDARY_RE.split(line.group(0)) + if part.startswith("uv sync") + ) + + +@pytest.mark.parametrize("relative_path", PROXY_DOCKERFILES) +def test_every_uv_sync_installs_bedrock_realtime_extra(relative_path: str): + dockerfile_path: Final = os.path.join(REPO_ROOT, relative_path) + if not os.path.exists(dockerfile_path): + pytest.skip(f"{relative_path} not present in this checkout") + + with open(dockerfile_path, "r", encoding="utf-8") as f: + contents: Final = f.read() + + invocations: Final = _uv_sync_invocations(contents) + assert invocations, f"{relative_path} has no `uv sync` invocation" + + missing: Final = tuple(invocation for invocation in invocations if "--extra bedrock-realtime" not in invocation) + assert not missing, ( + f"{relative_path}: {len(missing)} of {len(invocations)} `uv sync` invocations omit " + "`--extra bedrock-realtime`, so aws-sdk-bedrock-runtime is absent and Bedrock Nova Sonic " + "/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'" + ) diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index 0458af0da0e..a7a9e0fc37d 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -84,3 +84,41 @@ def test_bare_fireworks_ids_resolve_through_prefixed_entries(): assert info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) assert info["max_input_tokens"] == expected["max_input_tokens"] assert info["max_output_tokens"] == expected["max_output_tokens"] + + +TWIN_PINNED_PRICES = { + "deepseek-v4-flash-0731": { + "input_cost_per_token": 2.2e-07, + "cache_read_input_token_cost": 7e-09, + "output_cost_per_token": 6.6e-07, + }, +} + + +def test_deepseek_v4_flash_0731_twins_pin_published_pricing(model_data): + """Both 0731 entries carry the price published at docs.fireworks.ai/serverless/pricing.""" + for bare_suffix, expected in TWIN_PINNED_PRICES.items(): + for key in ( + f"fireworks_ai/{bare_suffix}", + f"fireworks_ai/accounts/fireworks/models/{bare_suffix}", + ): + entry = model_data[key] + for field, value in expected.items(): + assert entry[field] == pytest.approx(value), f"{key}.{field}" + + +def test_fireworks_account_prefixed_twins_agree_on_price(model_data): + """Every accounts/fireworks/models/X entry prices identically to its bare fireworks_ai/X twin.""" + prefix = "fireworks_ai/accounts/fireworks/models/" + pairs_checked = 0 + for key, entry in model_data.items(): + if not key.startswith(prefix): + continue + bare_key = f"fireworks_ai/{key[len(prefix):]}" + bare_entry = model_data.get(bare_key) + if bare_entry is None: + continue + pairs_checked += 1 + for field in sorted({f for f in (*entry, *bare_entry) if "cost" in f}): + assert entry.get(field) == bare_entry.get(field), f"{key} vs {bare_key}: {field}" + assert pairs_checked >= 20 diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..7c2b9d0be05 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3150,8 +3150,8 @@ def _stream_builder_logging_obj() -> LiteLLMLogging: return logging_obj -def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) +def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) chunks: Final = [ _stream_builder_text_chunk("gpt-4o", "Hello "), _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), @@ -3168,11 +3168,45 @@ def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypa assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) -def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) +def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj: Final = LiteLLMLogging( + model="us.anthropic.claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-unpriceable", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response._hidden_params.get("response_cost") is None + + +def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) chunks: Final = [ _stream_builder_text_chunk("gpt-4o", "Hello "), _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, ] response: Final = litellm.stream_chunk_builder( @@ -3180,4 +3214,26 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( ) assert response is not None - assert response._hidden_params.get("response_cost") is None + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + assert response._hidden_params["response_cost"] == pytest.approx(0.5) + + +def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): + from openai.types.completion_usage import CompletionUsage + + usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) + assert type(usage_chunk.usage) is CompletionUsage + chunks: Final = [ + _stream_builder_text_chunk("mantle-claude", "Hello "), + _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response.usage.prompt_tokens == 20 + assert response.usage.completion_tokens == 60 + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + assert response._hidden_params["response_cost"] == pytest.approx(0.000704) diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py index 94e4e3c81e5..7a42eaf0f0a 100644 --- a/tests/test_litellm/test_openai_embedding_encoding_format_default.py +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -1,124 +1,121 @@ -from unittest.mock import MagicMock, patch +import json +from typing import Final +import httpx import pytest +import respx -from litellm import embedding +import litellm -@pytest.mark.parametrize( - "set_env, env_value, expected", - [ - (False, None, "float"), - (True, "base64", "base64"), - ], -) -def test_openai_embedding_encoding_format_default( - monkeypatch, set_env, env_value, expected -): - monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - if set_env: - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) - - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } +def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) ) - mock_response.headers = {} - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) +@pytest.fixture(autouse=True) +def clear_default_encoding_format_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == expected + +def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_forwards_explicit_encoding_format(respx_mock: respx.MockRouter) -> None: + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +def test_embedding_openai_explicit_encoding_format_wins_over_env_var( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", encoding_format="base64" + ) + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == "base64" + + +@pytest.mark.parametrize("env_value", ["float", "base64"]) +def test_embedding_openai_env_var_sets_default_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_value: str +) -> None: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert request_body["encoding_format"] == env_value @pytest.mark.parametrize("env_none", ["none", "NONE", " none "]) -def test_openai_embedding_encoding_format_env_none_omits_param( - monkeypatch, env_none -): - """LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT=none omits encoding_format (provider default).""" +def test_embedding_openai_env_none_omits_encoding_format( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, env_none: str +) -> None: monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_none) + mock_route: Final = _mock_openai_embedding_route(respx_mock) - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } + litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + + +@pytest.mark.asyncio +async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + mock_route: Final = _mock_openai_embedding_route(respx_mock) + + response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + + request_body: Final = json.loads(mock_route.calls.last.request.read()) + assert "encoding_format" not in request_body + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_openai_omitted_encoding_format_maps_provider_errors( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + respx_mock.post("https://api.openai.com/v1/embeddings").mock( + return_value=httpx.Response( + 429, + headers={"retry-after": "42", "x-should-retry": "false"}, + json={"error": {"message": "rate limited", "type": "rate_limit_error"}}, + ) ) - mock_response.headers = {} - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response + with pytest.raises(litellm.RateLimitError) as exc_info: + litellm.embedding( + model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test", max_retries=0 ) - embedding( - model="text-embedding-ada-002", - input="Hello world", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert "encoding_format" not in call_kwargs - - -def test_openai_embedding_encoding_format_explicit_overrides_env(monkeypatch): - """Request `encoding_format` wins over LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT.""" - monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") - - mock_response = MagicMock() - mock_response.parse.return_value = MagicMock( - model_dump=lambda: { - "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], - "model": "text-embedding-ada-002", - "object": "list", - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - ) - mock_response.headers = {} - - with patch( - "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" - ) as mock_get_client: - mock_client_instance = MagicMock() - mock_get_client.return_value = mock_client_instance - mock_client_instance.embeddings.with_raw_response.create.return_value = ( - mock_response - ) - - embedding( - model="text-embedding-ada-002", - input="Hello world", - encoding_format="base64", - ) - - call_kwargs = ( - mock_client_instance.embeddings.with_raw_response.create.call_args[1] - ) - assert call_kwargs["encoding_format"] == "base64" + assert int(exc_info.value.litellm_response_headers["retry-after"]) == 42 diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py new file mode 100644 index 00000000000..c0860a5b55f --- /dev/null +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -0,0 +1,156 @@ +import json +from functools import lru_cache +from pathlib import Path + +import pytest + +import litellm + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +FLEX_LONG_CONTEXT = { + "gpt-5.4": { + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + }, + "gpt-5.4-pro": { + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + }, + "gpt-5.5": { + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + }, +} + +PRIORITY_LONG_CONTEXT = { + "gpt-5.6": { + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + }, + "gpt-5.6-sol": { + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + }, + "gpt-5.6-terra": { + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + }, + "gpt-5.6-luna": { + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + }, +} + +EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} + +NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") + + +@pytest.fixture(autouse=True) +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + +@lru_cache(maxsize=2) +def _load(path: Path) -> dict[str, dict[str, object]]: + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("path", [MAIN_PATH, BACKUP_PATH], ids=["main", "backup"]) +@pytest.mark.parametrize("model", sorted(EXPECTED)) +def test_service_tier_long_context_rates_are_published(model: str, path: Path) -> None: + """Each tier must carry its own above-272K rates, in both price files.""" + info = _load(path).get(model) + assert info is not None, f"{model} not found in {path.name}" + for key, expected in EXPECTED[model].items(): + assert info.get(key) == pytest.approx(expected), f"{model}.{key} is {info.get(key)!r}, expected {expected!r}" + + +@pytest.mark.parametrize("model", sorted(EXPECTED)) +def test_tier_long_context_rate_is_half_or_double_the_standard(model: str) -> None: + """Flex is half the standard long-context rate; priority is double it.""" + info = _load(MAIN_PATH)[model] + tier = "flex" if model in FLEX_LONG_CONTEXT else "priority" + ratio = 0.5 if tier == "flex" else 2.0 + for base in ("input_cost_per_token", "output_cost_per_token"): + standard = info[f"{base}_above_272k_tokens"] + tiered = info[f"{base}_above_272k_tokens_{tier}"] + assert tiered == pytest.approx(standard * ratio), ( + f"{model}.{base}_above_272k_tokens_{tier} is {tiered!r}, " + f"expected {ratio}x the standard long-context rate {standard!r}" + ) + + +@pytest.mark.parametrize("model", NO_PUBLISHED_PRIORITY_LONG_CONTEXT) +def test_no_priority_long_context_rates_where_openai_publishes_none(model: str) -> None: + """Guard against back-filling a rate OpenAI does not publish.""" + info = _load(MAIN_PATH)[model] + assert "input_cost_per_token_above_272k_tokens_priority" not in info + + +LONG_CONTEXT_PROMPT_TOKENS = 300_000 +COMPLETION_TOKENS = 1_000 + +TIERED_COST_CASES = [ + ("gpt-5.4", "flex", 2.5e-06, 1.125e-05), + ("gpt-5.4-pro", "flex", 3e-05, 0.000135), + ("gpt-5.5", "flex", 5e-06, 2.25e-05), + ("gpt-5.6", "priority", 1.6e-05, 6e-05), + ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05), + ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), + ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), +] + + +@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES) +def test_cost_per_token_bills_long_context_at_the_tier_rate( + model: str, tier: str, input_rate: float, output_rate: float +) -> None: + """A prompt over 272K on flex or priority must bill at that tier's long-context rate.""" + input_cost, output_cost = litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + service_tier=tier, + ) + assert input_cost == pytest.approx(LONG_CONTEXT_PROMPT_TOKENS * input_rate) + assert output_cost == pytest.approx(COMPLETION_TOKENS * output_rate) + + +@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES) +def test_cost_per_token_tier_differs_from_the_standard_long_context_cost( + model: str, tier: str, input_rate: float, output_rate: float +) -> None: + """Flex halves the standard long-context bill and priority doubles it.""" + ratio = 0.5 if tier == "flex" else 2.0 + standard = sum( + litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + ) + ) + tiered = sum( + litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + service_tier=tier, + ) + ) + assert tiered == pytest.approx(standard * ratio) diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 39f498b4e58..452a15334ef 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -793,3 +793,203 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): finally: litellm.model_cost.pop(model_key, None) _invalidate_model_cost_lowercase_map() + + +def test_update_dictionary_merges_nested_dicts_without_aliasing(): + """A nested dict must be merged copy-on-write: the pre-existing nested dict + object stays untouched, and the caller's incoming nested dict is never + inserted by reference into the merged result. + """ + from litellm.utils import _update_dictionary + + existing_nested = {"hours_utc": "01:00-02:00"} + existing = {"off_peak_pricing": existing_nested} + incoming_nested = {"windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}]} + incoming = {"off_peak_pricing": incoming_nested} + + merged = _update_dictionary(existing, incoming) + + assert merged["off_peak_pricing"] == { + "hours_utc": "01:00-02:00", + "windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}], + } + assert existing_nested == {"hours_utc": "01:00-02:00"} + assert merged["off_peak_pricing"] is not incoming_nested + + fresh = _update_dictionary({}, incoming) + assert fresh["off_peak_pricing"] == incoming_nested + assert fresh["off_peak_pricing"] is not incoming_nested + + +def test_router_deployments_sharing_backend_keep_their_own_off_peak_pricing(): + """Two deployments of the same backend model with different + ``off_peak_pricing`` blocks must each keep their own schedule under their + unique model id, and neither block may leak onto the shared backend keys. + + Before the fix, ``register_model`` inserted the first deployment's block by + reference into the built-in ``gpt-4o-mini`` entry, and the second + deployment's registration merged its keys into that same object, corrupting + the first deployment's schedule and polluting the built-in entry. + """ + from litellm import Router + + active_block = { + "windows": [{"hours_utc": "16:00-19:00", "weekdays": [2]}], + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1e-06, + } + inactive_block = { + "hours_utc": "05:00-06:00", + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1e-06, + } + shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"] + deployment_ids = ["offpeak-alias-dep-1", "offpeak-alias-dep-2"] + original_entries = _snapshot_model_cost_entries(shared_keys) + + router = Router( + model_list=[ + { + "model_name": "offpeak-active-weekday", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key-for-registration", + }, + "model_info": { + "id": deployment_ids[0], + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "off_peak_pricing": dict(active_block), + }, + }, + { + "model_name": "offpeak-inactive-hours", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key-for-registration", + }, + "model_info": { + "id": deployment_ids[1], + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "off_peak_pricing": dict(inactive_block), + }, + }, + ] + ) + + try: + registered_first = litellm.model_cost[deployment_ids[0]]["off_peak_pricing"] + registered_second = litellm.model_cost[deployment_ids[1]]["off_peak_pricing"] + assert registered_first == active_block + assert registered_second == inactive_block + for shared_key in shared_keys: + shared_entry = litellm.model_cost.get(shared_key) or {} + assert not shared_entry.get("off_peak_pricing") + finally: + for deployment_id in deployment_ids: + litellm.model_cost.pop(deployment_id, None) + _restore_model_cost_entries(original_entries) + del router + + +def test_router_off_peak_only_deployment_inherits_builtin_base_rates(): + """A deployment that sets only ``off_peak_pricing`` on its model_info must + still be costed from its deployment-scoped entry: the base token rates are + inherited from the backend model's built-in cost map entry, since the + shared backend key deliberately never carries the off-peak block. + """ + from litellm import Router + + block = { + "hours_utc": "00:00-00:00", + "input_cost_per_token": 5e-05, + "output_cost_per_token": 1e-04, + } + shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"] + deployment_id = "offpeak-only-dep-1" + original_entries = _snapshot_model_cost_entries(shared_keys + [deployment_id]) + builtin_info = litellm.get_model_info(model="openai/gpt-4o-mini") + + router = Router( + model_list=[ + { + "model_name": "offpeak-only", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key-for-registration", + }, + "model_info": {"id": deployment_id, "off_peak_pricing": dict(block)}, + } + ] + ) + + try: + entry = litellm.model_cost[deployment_id] + assert entry["off_peak_pricing"] == block + assert entry["input_cost_per_token"] is not None + assert entry["input_cost_per_token"] == builtin_info["input_cost_per_token"] + assert entry["output_cost_per_token"] == builtin_info["output_cost_per_token"] + for shared_key in shared_keys: + shared_entry = litellm.model_cost.get(shared_key) or {} + assert not shared_entry.get("off_peak_pricing") + finally: + _restore_model_cost_entries(original_entries) + del router + + +def test_use_custom_pricing_for_model_sees_off_peak_only_model_info(): + from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model + + block = {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-05} + assert use_custom_pricing_for_model({"metadata": {"model_info": {"off_peak_pricing": block}}}) is True + assert use_custom_pricing_for_model({"metadata": {"model_info": {"off_peak_pricing": None}}}) is False + assert use_custom_pricing_for_model({"metadata": {"model_info": {"id": "some-id"}}}) is False + + +def test_completion_cost_applies_off_peak_only_deployment_pricing(): + """End to end through the cost calculator: with ``custom_pricing`` set and + a ``router_model_id`` whose entry carries only an always-on off-peak block, + the request bills at the block's rates rather than the shared backend rate. + """ + from litellm import Router + from litellm.types.utils import ModelResponse, Usage + + block = { + "hours_utc": "00:00-00:00", + "input_cost_per_token": 5e-05, + "output_cost_per_token": 1e-04, + } + shared_keys = ["gpt-4o-mini", "openai/gpt-4o-mini"] + deployment_id = "offpeak-only-dep-2" + original_entries = _snapshot_model_cost_entries(shared_keys + [deployment_id]) + + router = Router( + model_list=[ + { + "model_name": "offpeak-only", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key-for-registration", + }, + "model_info": {"id": deployment_id, "off_peak_pricing": dict(block)}, + } + ] + ) + + try: + response = ModelResponse( + model="gpt-4o-mini", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + cost = litellm.completion_cost( + completion_response=response, + model="openai/gpt-4o-mini", + custom_llm_provider="openai", + custom_pricing=True, + router_model_id=deployment_id, + ) + assert cost == pytest.approx(100 * 5e-05 + 50 * 1e-04) + finally: + _restore_model_cost_entries(original_entries) + del router diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 44c1cdbff06..d6328118f57 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7271,6 +7271,71 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_configured_display_name_reads_deployment_model_info(): + router = litellm.Router( + model_list=[ + { + "model_name": "Kimi K3-claude-compatible", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"display_name": "Kimi K3"}, + } + ] + ) + + assert router.get_configured_display_name("Kimi K3-claude-compatible") == "Kimi K3" + + +def test_get_configured_display_name_returns_none_for_unset_or_unknown(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-display-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_configured_display_name("no-display-model") is None + assert router.get_configured_display_name("not-a-real-model") is None + + +def test_get_configured_display_name_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"display_name": "Bedrock"}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + +def test_get_configured_display_name_treats_malformed_values_as_absent(): + malformed = ["", " ", 12345, ["Kimi K3"], {"name": "Kimi K3"}, True] + router = litellm.Router( + model_list=[ + { + "model_name": f"bad-display-{i}", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"display_name": bad}, + } + for i, bad in enumerate(malformed) + ] + ) + + for i in range(len(malformed)): + assert router.get_configured_display_name(f"bad-display-{i}") is None + + @pytest.mark.asyncio async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): router = litellm.Router( @@ -11595,3 +11660,138 @@ class TestTierParamsTheTargetAccepts: accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) assert accepted == {"reasoning_effort": "max"} + + +class TestPreRoutingTierDrivesFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup stayed on the router name, so the tier's configured chain never ran and a + provider failure on the tier's first hop was returned to the client.""" + + class _TierRouter(litellm.Router): + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + if model == "smart-router": + return PreRoutingHookResponse(model="tier1", messages=messages) + return None + + @classmethod + def _router(cls, fallbacks) -> "litellm.Router": + return cls._TierRouter( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, + }, + { + "model_name": "tier1", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "backup-a", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-a", + }, + }, + { + "model_name": "backup-b", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-b", + }, + }, + { + "model_name": "failing-backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "plain", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + ], + fallbacks=fallbacks, + num_retries=0, + ) + + @pytest.mark.asyncio + async def test_the_selected_tier_fallback_chain_runs(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_is_not_used(self): + """The router name has no chain of its own, so nothing should rescue this call.""" + router = self._router([{"tier2": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self): + """The documented contract: configs keyed on the requested name keep working behind auto-routers.""" + router = self._router([{"smart-router": ["backup-a"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_the_tier_chain_wins_over_the_router_name_chain(self): + router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="tier1", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self): + """The metadata bucket carries caller-supplied keys, so only the hook may set the tier.""" + router = self._router([{"tier1": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion( + model="plain", + messages=[{"role": "user", "content": "hi"}], + metadata={"pre_routing_selected_model": "tier1"}, + ) + + @pytest.mark.asyncio + async def test_each_fallback_hop_resolves_its_own_chain(self): + """The second hop must key off the group it is running, not the tier that failed.""" + router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-b" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index b580b03574e..30b265905f3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -538,6 +538,157 @@ def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): assert model_info == {"input_cost_per_token": 0.000003} +def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates(): + """Direct unit test of the helper: an entry carrying only an + off_peak_pricing block inherits the backend model's built-in base token + rates, so cost lookup via the deployment id can bill standard rates + outside the windows. + """ + backend_model = "gpt-4o-mini" + builtin_info = litellm.get_model_info(model=backend_model, custom_llm_provider="openai") + off_peak_block = { + "hours_utc": "00:00-00:00", + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1e-06, + } + model_info = {"off_peak_pricing": off_peak_block} + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="openai", + ) + + assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] + assert model_info["output_cost_per_token"] == builtin_info["output_cost_per_token"] + assert model_info["off_peak_pricing"] == off_peak_block + + +def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): + """A backend with above-threshold pricing hands the whole rate structure to + the deployment entry, so peak-hour billing of large prompts through that + entry matches the shared backend entry instead of flattening to the base + rate. + """ + backend_model = "gemini/gemini-2.5-pro" + builtin_info = litellm.get_model_info(model=backend_model) + assert builtin_info["input_cost_per_token_above_200k_tokens"] is not None + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="gemini", + ) + + assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] + assert ( + model_info["input_cost_per_token_above_200k_tokens"] + == builtin_info["input_cost_per_token_above_200k_tokens"] + ) + assert ( + model_info["output_cost_per_token_above_200k_tokens"] + == builtin_info["output_cost_per_token_above_200k_tokens"] + ) + + +def test_inherit_builtin_base_rates_for_off_peak_carries_companion_billing_fields(): + """Billing rules that are not literal cost rates, like the web search + billing unit, must ride along, or grounding and regional uplifts would + bill differently through the deployment entry than through the shared + backend entry. + """ + backend_model = "gemini-3-pro-image" + raw_entry = litellm.model_cost[backend_model] + assert raw_entry.get("web_search_billing_unit") is not None + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider=None, + ) + + assert model_info["web_search_billing_unit"] == raw_entry["web_search_billing_unit"] + assert model_info["input_cost_per_token"] == raw_entry["input_cost_per_token"] + + +def test_inherit_builtin_base_rates_for_off_peak_tiered_only_backend_stores_no_zero(): + """A tiered-only backend has no flat token rates; get_model_info synthesizes + zeros for them, and storing those would mark the deployment explicitly + priced free. The tier table itself must carry over as an isolated copy so + mutating the deployment entry never touches the shared cost map. + """ + backend_model = "dashscope/qwen-flash" + raw_tiers = litellm.model_cost[backend_model]["tiered_pricing"] + + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="dashscope", + ) + + assert model_info.get("input_cost_per_token") != 0 + assert model_info.get("output_cost_per_token") != 0 + assert model_info["tiered_pricing"] == raw_tiers + assert model_info["tiered_pricing"] is not raw_tiers + assert model_info["tiered_pricing"][0] is not raw_tiers[0] + + original_first_tier = copy.deepcopy(raw_tiers[0]) + model_info["tiered_pricing"][0]["input_cost_per_token"] = 123.0 + assert raw_tiers[0] == original_first_tier + + +def test_inherit_builtin_base_rates_for_off_peak_leaves_explicit_rates_alone(): + """An entry that sets its own base rate beside the block already counts as + a full custom pricing entry; the helper must not mix builtin rates into it. + """ + model_info = { + "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, + "input_cost_per_token": 3e-06, + } + + Router._inherit_builtin_base_rates_for_off_peak( + model_info=model_info, + backend_model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + assert model_info["input_cost_per_token"] == 3e-06 + assert "output_cost_per_token" not in model_info + + +def test_inherit_builtin_base_rates_for_off_peak_noop_without_block_or_backend(): + """Nothing happens without an off_peak_pricing block, and an unmapped + backend model leaves the entry unchanged rather than raising. + """ + plain_info = {"id": "dep-1"} + Router._inherit_builtin_base_rates_for_off_peak( + model_info=plain_info, + backend_model="gpt-4o-mini", + custom_llm_provider="openai", + ) + assert plain_info == {"id": "dep-1"} + + off_peak_info = {"off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}} + Router._inherit_builtin_base_rates_for_off_peak( + model_info=off_peak_info, + backend_model="this-backend-model-does-not-exist-x9y8z7", + custom_llm_provider=None, + ) + assert "input_cost_per_token" not in off_peak_info + + def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields(): """The shared-backend-key stripping in Router relies on CustomPricingLiteLLMParams enumerating every per-deployment pricing field. @@ -2130,3 +2281,83 @@ def test_every_declaring_deployment_is_named(caplog): assert "azure-ptu-east" in warnings[0] assert "azure-ptu-west" in warnings[0] assert "plain-gpt-4o" not in warnings[0] + + +def _simulate_price_data_reload_with_provider_sets(monkeypatch, fetched_catalog): + """Like `_simulate_price_data_reload`, plus the provider model-set refresh the proxy's + `_swap_in_model_cost_map` does before replaying, so bare names in the new catalog resolve.""" + monkeypatch.setattr(litellm, "model_cost", fetched_catalog) + _invalidate_model_cost_lowercase_map() + litellm.add_known_models(model_cost_map=fetched_catalog) + reapply_runtime_model_cost_registrations() + + +def test_a_config_deployment_dropped_by_a_stale_cost_map_comes_back_on_reload(monkeypatch): + """ + Booting on the bundled backup, a bare model that only the remote catalog knows + cannot be provider-resolved, so the proxy router (ignore_invalid_deployments) drops + it. Once a reload brings in a catalog that knows the model, the deployment must be + served again with its access groups, and exactly once however many reloads follow. + """ + backend = "lit-5766-only-in-remote-catalog" + try: + router = Router( + model_list=[ + { + "model_name": "new-model", + "litellm_params": {"model": backend, "api_key": "k"}, + "model_info": {"id": "new-id", "access_groups": ["team-models"]}, + }, + { + "model_name": "control-model", + "litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"}, + "model_info": {"id": "control-id", "access_groups": ["team-models"]}, + }, + ], + ignore_invalid_deployments=True, + ) + assert router.get_model_names() == ["control-model"] + assert router.get_model_access_groups(model_name="new-model") == {} + + fresh_catalog = {**litellm.model_cost, backend: {"litellm_provider": "openai", "mode": "chat"}} + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + + assert sorted(router.get_model_names()) == ["control-model", "new-model"] + assert router.get_model_access_groups(model_name="new-model") == {"team-models": ["new-model"]} + assert [d["model_info"]["id"] for d in router.model_list] == ["control-id", "new-id"] + assert "new-id" in litellm.model_cost + finally: + litellm.open_ai_chat_completion_models.discard(backend) + litellm.models_by_provider["openai"].discard(backend) + + +def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_reload(monkeypatch): + """ + Only provider-resolution drops can be healed by a fresh catalog. A deployment that + fails after its provider resolved (here a pass-through vertex entry with no project) + has already touched router state, so replaying it on every reload would leak into + `deployment_names` each time. + """ + router = Router( + model_list=[ + { + "model_name": "vertex-passthrough", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "use_in_pass_through": True}, + "model_info": {"id": "vertex-id"}, + }, + { + "model_name": "control-model", + "litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"}, + "model_info": {"id": "control-id"}, + }, + ], + ignore_invalid_deployments=True, + ) + assert router.get_model_names() == ["control-model"] + names_after_boot = list(router.deployment_names) + + _simulate_price_data_reload_with_provider_sets(monkeypatch, dict(litellm.model_cost)) + + assert router.get_model_names() == ["control-model"] + assert router.deployment_names == names_after_boot diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/test_litellm/test_router_order_fallback.py index 7743cb005d0..fde870e5abe 100644 --- a/tests/test_litellm/test_router_order_fallback.py +++ b/tests/test_litellm/test_router_order_fallback.py @@ -6,12 +6,19 @@ should be tried first, and higher order deployments should be used as fallbacks when lower order deployments fail. """ -from typing import Optional +import json +from typing import Final, Optional +import httpx import pytest +from openai import AsyncOpenAI +import litellm from litellm import Router -from litellm.utils import _get_order_filtered_deployments +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.prompt_caching_cache import PromptCachingCache +from litellm.types.router import RouterRateLimitError +from litellm.utils import _get_deployment_order, _get_order_filtered_deployments # --------------------------------------------------------------------------- # Unit tests for _get_order_filtered_deployments @@ -49,13 +56,22 @@ class TestGetOrderFilteredDeployments: assert len(result) == 1 assert result[0]["model_info"]["id"] == "b" - def test_target_order_no_match_returns_all(self): + def test_target_order_no_match_returns_empty(self): deps = [ self._make_deployment(1, "a"), self._make_deployment(2, "b"), ] result = _get_order_filtered_deployments(deps, target_order=99) - assert len(result) == 2 + assert result == [] + + def test_target_order_no_match_does_not_reselect_lower_order(self): + deps = [ + self._make_deployment(1, "a"), + self._make_deployment(2, "b"), + ] + remaining_after_pre_call = [deps[0]] + result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2) + assert result == [] def test_no_order_set_returns_all(self): deps = [ @@ -406,35 +422,239 @@ async def test_router_order_fallback_with_hidden_model_group_alias(): assert response._hidden_params["model_id"] == "2" +@pytest.mark.asyncio +async def test_router_order_fallback_does_not_reselect_order_1_when_order_2_is_filtered_out(): + class _DropOrder2(CustomLogger): + async def async_filter_deployments( + self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None + ): + return [d for d in healthy_deployments if _get_deployment_order(d) != 2] + + drop_order_2: Final = _DropOrder2() + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "key", + "mock_response": "litellm.RateLimitError", + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "key", + "mock_response": "success from order 2", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + ) + litellm.callbacks.append(drop_order_2) + try: + with pytest.raises(RouterRateLimitError, match="No deployments available") as exc_info: + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert "success from order 2" not in str(exc_info.value) + finally: + litellm.callbacks.remove(drop_order_2) + + +@pytest.mark.asyncio +async def test_router_order_fallback_ignores_prompt_cache_pin_on_target_order(): + messages = [{"role": "user", "content": "word " * 5000}] + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("azure peak load"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "success from order 2", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + optional_pre_call_checks=["prompt_caching"], + ) + await PromptCachingCache(cache=router.cache).async_add_model_id( + model_id="1", + messages=messages, + tools=None, + ) + response = await router.acompletion(model="test-model", messages=messages) + assert response._hidden_params["model_id"] == "2" + + +@pytest.mark.asyncio +async def test_router_order_fallback_retries_keep_target_order(): + seen_target_orders: Final = [] + + class _RecordTargetOrder(CustomLogger): + async def async_filter_deployments( + self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None + ): + seen_target_orders.append((request_kwargs or {}).get("_target_order")) + return healthy_deployments + + recorder: Final = _RecordTargetOrder() + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("fail order 1"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("fail order 2"), + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=1, + ) + litellm.callbacks.append(recorder) + try: + with pytest.raises(Exception, match="fail order 2"): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + finally: + litellm.callbacks.remove(recorder) + assert seen_target_orders.count(2) >= 2 + + +@pytest.mark.asyncio +async def test_generic_api_call_strips_target_order_from_provider_kwargs(): + captured: Final = {} + + async def _fake_provider(**provider_kwargs): + captured.update(provider_kwargs) + return "ok" + + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key", "order": 2}, + "model_info": {"id": "2"}, + }, + ], + ) + response = await router._ageneric_api_call_with_fallbacks_helper( + model="test-model", + original_generic_function=_fake_provider, + _target_order=2, + messages=[{"role": "user", "content": "hi"}], + ) + assert response == "ok" + assert captured["model"] == "gpt-4o" + assert "_target_order" not in captured + + +@pytest.mark.asyncio +async def test_text_completion_order_fallback_hop_does_not_send_target_order_upstream(): + upstream_bodies: Final[list[dict]] = [] + + def _upstream(request: httpx.Request) -> httpx.Response: + upstream_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "id": "cmpl-1", + "object": "text_completion", + "created": 0, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "ok from order 2", "index": 0, "logprobs": None, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + upstream_client: Final = AsyncOpenAI( + api_key="key", + base_url="http://upstream.test", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_upstream)), + ) + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_key": "key", + "mock_response": Exception("fail order 1"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_key": "key", + "api_base": "http://upstream.test", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + ) + try: + response = await router.atext_completion(model="test-model", prompt="hi", client=upstream_client) + finally: + await upstream_client.close() + + assert response._hidden_params["model_id"] == "2" + assert upstream_bodies + assert all("_target_order" not in body for body in upstream_bodies) + + def test_check_non_standard_fallback_format(): from litellm.router_utils.fallback_event_handlers import ( _check_non_standard_fallback_format, ) # Standard formats - assert ( - _check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) - == False - ) + assert _check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) == False assert _check_non_standard_fallback_format([{"model": ["qwen-backup"]}]) == False - assert ( - _check_non_standard_fallback_format( - [{"model": ["qwen-backup"], "region": ["us-east-1"]}] - ) - == False - ) + assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "region": ["us-east-1"]}]) == False # Non-standard formats assert _check_non_standard_fallback_format([{"model": "qwen-backup"}]) == True assert ( - _check_non_standard_fallback_format( - [{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}] - ) - == True - ) - assert ( - _check_non_standard_fallback_format( - [{"model": ["qwen-backup"], "api_key": "some-key"}] - ) + _check_non_standard_fallback_format([{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}]) == True ) + assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "api_key": "some-key"}]) == True diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index a7fa6ac5568..0790b41c349 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1005,6 +1005,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "gemini_native_audio": {"type": "boolean"}, "gemini_audio_only_live": {"type": "boolean"}, "supports_embedding_image_input": {"type": "boolean"}, + "supports_forced_tool_use": {"type": "boolean"}, "supports_function_calling": {"type": "boolean"}, "supports_image_input": {"type": "boolean"}, "supports_nova_canvas_image_edit": {"type": "boolean"}, @@ -1416,6 +1417,26 @@ def test_get_provider_rerank_config(): assert isinstance(config, HostedVLLMRerankConfig) +def test_get_provider_text_to_speech_config_vertex_gemini_skips_cloud_tts(): + """Regression for LIT-6501: mapping vertex Gemini TTS params through Google Cloud TTS + dropped response_format before the speech_to_completion bridge could honor it.""" + from litellm.llms.vertex_ai.text_to_speech.transformation import VertexAITextToSpeechConfig + from litellm.utils import LlmProviders + + assert ( + ProviderConfigManager.get_provider_text_to_speech_config( + model="gemini-2.5-flash-preview-tts", provider=LlmProviders.VERTEX_AI + ) + is None + ) + assert isinstance( + ProviderConfigManager.get_provider_text_to_speech_config( + model="en-US-Studio-O", provider=LlmProviders.VERTEX_AI + ), + VertexAITextToSpeechConfig, + ) + + # Models that should be skipped during testing OLD_PROVIDERS = ["aleph_alpha", "palm"] SKIP_MODELS = [ @@ -4634,6 +4655,7 @@ GEMINI_4096_CACHE_MIN_MODELS: Final = tuple( "gemini-3.5-flash", "gemini-3.6-flash", "gemini-3.7-flash", + "gemini-3.8-flash", "gemini-3.1-pro-preview", "gemini-3.1-pro-preview-customtools", ) diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index b166e902d6e..2a60ff9c4b5 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -532,6 +532,41 @@ class TestVideoGeneration: assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_completion_cost_veo_31_tiers_pin_published_rates(self, monkeypatch): + """The gemini and vertex_ai veo 3.1 entries bill Google's published per-second tier rates.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, provider: str, resolution: str | None, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": duration, + **({"video_resolution": resolution} if resolution else {}), + } + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider=provider, + ) + + for provider in ("gemini", "vertex_ai"): + for suffix in ("generate-preview", "generate-001"): + standard = f"{provider}/veo-3.1-{suffix}" + fast = f"{provider}/veo-3.1-fast-{suffix}" + assert abs(cost_for(standard, provider, None, 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "1080p", 8.0) - 3.2) < 1e-6 + assert abs(cost_for(standard, provider, "4k", 8.0) - 4.8) < 1e-6 + assert abs(cost_for(fast, provider, "720p", 8.0) - 0.8) < 1e-6 + assert abs(cost_for(fast, provider, "1080p", 8.0) - 0.96) < 1e-6 + assert abs(cost_for(fast, provider, "4k", 8.0) - 2.4) < 1e-6 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6eb7e115586..52cb9628252 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22519 + "limit": 22364 }, "LIT002": { - "limit": 26818 + "limit": 26777 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16544 + "limit": 16507 }, "LIT011": { - "limit": 5575 + "limit": 5535 }, "LIT012": { "limit": 4495 diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index bbf69c4a77a..e8207d179bd 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -4,5 +4,8 @@ "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, "local/no-large-inline-object-arg": { "max": 559, "target": 300 }, - "local/no-long-condition-chain": { "max": 265, "target": 120 } + "local/no-long-condition-chain": { "max": 265, "target": 120 }, + "testing-library/no-container": { "max": 133, "target": 50 }, + "testing-library/no-node-access": { "max": 716, "target": 500 }, + "testing-library/prefer-screen-queries": { "max": 18, "target": 18 } } diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index 23cc5096bb1..f5e3b23b3ec 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -104,10 +104,13 @@ const eslintConfig = [ plugins: { "testing-library": testingLibrary, "jest-dom": jestDom }, rules: { "testing-library/await-async-queries": "error", + "testing-library/no-container": "warn", + "testing-library/no-node-access": "warn", "testing-library/no-wait-for-multiple-assertions": "error", "testing-library/no-wait-for-side-effects": "error", "testing-library/prefer-find-by": "error", "testing-library/prefer-presence-queries": "error", + "testing-library/prefer-screen-queries": "warn", "jest-dom/prefer-checked": "error", "jest-dom/prefer-empty": "error", "jest-dom/prefer-enabled-disabled": "error", diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index d2a64b93384..4e5b0c1dda6 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -4891,9 +4891,9 @@ } }, "node_modules/baseline-browser-mapping": { - "version": "2.10.27", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.10.27.tgz", - "integrity": "sha512-zEs/ufmZoUd7WftKpKyXaT6RFxpQ5Qm9xytKRHvJfxFV9DFJkZph9RvJ1LcOUi0Z1ZVijMte65JbILeV+8QQEA==", + "version": "2.11.20", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz", + "integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==", "license": "Apache-2.0", "bin": { "baseline-browser-mapping": "dist/cli.cjs" @@ -4939,9 +4939,9 @@ } }, "node_modules/browserslist": { - "version": "4.28.2", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.2.tgz", - "integrity": "sha512-48xSriZYYg+8qXna9kwqjIVzuQxi+KYWp2+5nCYnYKPTr0LvD89Jqk2Or5ogxz0NUMfIjhh2lIUX/LyX9B4oIg==", + "version": "4.28.8", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz", + "integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==", "dev": true, "funding": [ { @@ -4959,11 +4959,11 @@ ], "license": "MIT", "dependencies": { - "baseline-browser-mapping": "^2.10.12", - "caniuse-lite": "^1.0.30001782", - "electron-to-chromium": "^1.5.328", - "node-releases": "^2.0.36", - "update-browserslist-db": "^1.2.3" + "baseline-browser-mapping": "^2.11.12", + "caniuse-lite": "^1.0.30001809", + "electron-to-chromium": "^1.5.402", + "node-releases": "^2.0.53", + "update-browserslist-db": "^1.3.0" }, "bin": { "browserslist": "cli.js" @@ -5042,9 +5042,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001791", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001791.tgz", - "integrity": "sha512-yk0l/YSrOnFZk3UROpDLQD9+kC1l4meK/wed583AXrzoarMGJcbRi2Q4RaUYbKxYAsZ8sWmaSa/DsLmdBeI1vQ==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "funding": [ { "type": "opencollective", @@ -5706,9 +5706,9 @@ } }, "node_modules/electron-to-chromium": { - "version": "1.5.349", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.349.tgz", - "integrity": "sha512-QsWVGyRuY07Aqb234QytTfwd5d9AJlfNIQ5wIOl1L+PZDzI9d9+Fn0FRale/QYlFxt/bUnB0/nLd1jFPGxGK1A==", + "version": "1.5.416", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.416.tgz", + "integrity": "sha512-K6bvB2BjnNrugtIih6ewlbBI9DXa976jIdiIlRLHhBoEI9a4JaQjjHyF+A1IQI543aQYR4LnmOrT/K5fZj0aPA==", "dev": true, "license": "ISC" }, @@ -9901,11 +9901,14 @@ } }, "node_modules/node-releases": { - "version": "2.0.38", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.38.tgz", - "integrity": "sha512-3qT/88Y3FbH/Kx4szpQQ4HzUbVrHPKTLVpVocKiLfoYvw9XSGOX2FmD2d6DrXbVYyAQTF2HeF6My8jmzx7/CRw==", + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", "dev": true, - "license": "MIT" + "license": "MIT", + "engines": { + "node": ">=18" + } }, "node_modules/nuqs": { "version": "2.9.4", @@ -12350,9 +12353,9 @@ } }, "node_modules/update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "funding": [ { diff --git a/ui/litellm-dashboard/public/assets/logos/alice.svg b/ui/litellm-dashboard/public/assets/logos/alice.svg new file mode 100644 index 00000000000..f18f887b98c --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/alice.svg @@ -0,0 +1,4 @@ + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx index 06099a9fc22..4d18ec2ef5f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx @@ -74,6 +74,42 @@ describe("AgentsTable", () => { expect(onDeleteClick).toHaveBeenCalledWith("agent-9", "Doomed Agent"); }); + it("filters agents by name or by agent card description", async () => { + const user = userEvent.setup(); + render( + , + ); + + const search = screen.getByPlaceholderText("Search agent names or descriptions..."); + await user.type(search, "billing"); + expect(screen.getByText("Billing Router")).toBeInTheDocument(); + expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); + + await user.clear(search); + await user.type(search, "support tickets"); + expect(screen.getByText("Second Agent")).toBeInTheDocument(); + expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); + }); + + it("shows the no-match empty state when the search matches nothing", async () => { + const user = userEvent.setup(); + render(); + + await user.type(screen.getByPlaceholderText("Search agent names or descriptions..."), "zzzz"); + expect(screen.queryByText("Test Agent")).not.toBeInTheDocument(); + expect(screen.getByText("No matching agents")).toBeInTheDocument(); + }); + it("hides the actions column entirely for non-admins", () => { const agent = makeAgent({ agent_id: "agent-2" }); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index 67c7ed74180..35ed6b66425 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -1,13 +1,15 @@ "use client"; import { SortingState } from "@tanstack/react-table"; -import { Bot, CircleCheck } from "lucide-react"; +import { Bot, CircleCheck, Search as SearchIcon, X } from "lucide-react"; import React, { useMemo, useState } from "react"; import { Agent } from "@/components/agents/types"; import { DataTable } from "@/components/shared/DataTable"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { Switch } from "@/components/ui/switch"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { filterBySearchTerm } from "@/utils/searchUtils"; import { getAgentsTableColumns } from "./AgentsTableColumns"; @@ -24,14 +26,18 @@ interface AgentsTableProps { const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; -function EmptyState() { +function EmptyState({ isFiltered }: { isFiltered: boolean }) { return (
-
No agents yet
-
Add an agent to make it available in your organization.
+
{isFiltered ? "No matching agents" : "No agents yet"}
+
+ {isFiltered + ? "Adjust the search to see more agents." + : "Add an agent to make it available in your organization."} +
); } @@ -47,6 +53,11 @@ const AgentsTable: React.FC = ({ onDeleteClick, }) => { const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [searchTerm, setSearchTerm] = useState(""); + const filteredAgents = useMemo( + () => filterBySearchTerm(agents, searchTerm, (agent) => [agent.agent_name, agent.agent_card_params?.description]), + [agents, searchTerm], + ); const columns = useMemo( () => getAgentsTableColumns({ isAdmin, onAgentClick, onDeleteClick }), @@ -55,7 +66,7 @@ const AgentsTable: React.FC = ({ return ( agent.agent_id || String(index)} sortingMode="client" @@ -63,10 +74,27 @@ const AgentsTable: React.FC = ({ onSortingChange={setSorting} isLoading={isLoading} loadingMessage="Loading agents…" - noDataMessage={} + noDataMessage={ 0} />} size="compact" toolbar={() => ( -
+
+ + + + + setSearchTerm(e.target.value)} + /> + {searchTerm && ( + + setSearchTerm("")}> + + + + )} + ({ createAgentCall: vi.fn(), @@ -309,8 +310,7 @@ describe("AddAgentForm submit payload", () => { await user.type(await screen.findByLabelText("Allowed Models"), "gpt-4o,"); await user.keyboard("{Escape}"); - await user.click(screen.getByLabelText("Allowed Agents (Sub-Agents)")); - await user.click(await screen.findByTitle("Sub Agent One")); + await chooseSelectOption(user, screen.getByLabelText("Allowed Agents (Sub-Agents)"), "Sub Agent One"); await user.keyboard("{Escape}"); await user.click(screen.getByText(/Configure which models, agents, and MCP tools/)); await user.click(screen.getByRole("button", { name: /^Next/ })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx index b2e0a42eecb..dae19032e20 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.test.tsx @@ -13,17 +13,17 @@ describe("APIReferenceView", () => { it("uses the API doc base url when provided", () => { const apiDocUrl = "https://docs.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(apiDocUrl)); }); it("falls back to the proxy base url when the docs url is missing", () => { const proxyUrl = "https://proxy.litellm.test"; - const { getAllByTestId } = render(); + render(); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); expect(codeBlocks[0]).toHaveTextContent(new RegExp(proxyUrl)); }); @@ -31,7 +31,7 @@ describe("APIReferenceView", () => { const apiDocUrl = "https://docs-preferred.litellm.test"; const proxyUrl = "https://proxy-backup.litellm.test"; - const { getAllByTestId } = render( + render( { />, ); - const codeBlocks = getAllByTestId(codeBlockTestId); + const codeBlocks = screen.getAllByTestId(codeBlockTestId); const renderedCode = codeBlocks[0].textContent ?? ""; expect(renderedCode).toContain(apiDocUrl); expect(renderedCode).not.toContain(proxyUrl); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx index bc920e55abd..17a297d203d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx @@ -4,6 +4,7 @@ import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import BudgetModal from "./budget_modal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); @@ -63,8 +64,7 @@ describe("BudgetModal", () => { await openOptionalSettings(user); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await create(user); @@ -80,8 +80,7 @@ describe("BudgetModal", () => { await openOptionalSettings(user); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await user.click(screen.getByText("Optional Settings")); await waitFor(() => expect(screen.queryByLabelText("Max Budget (USD)")).not.toBeInTheDocument()); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx index 3fa96b54f1d..fe0ecfc10dd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx @@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { components } from "@/lib/http/schema"; import EditBudgetModal from "./edit_budget_modal"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const { updateMock } = vi.hoisted(() => ({ updateMock: vi.fn() })); @@ -73,8 +74,7 @@ describe("EditBudgetModal", () => { await user.clear(screen.getByLabelText("Max Budget (USD)")); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("monthly")); + await chooseSelectOption(user, screen.getByRole("combobox"), "monthly"); await save(user); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx index 9d4d5a6d425..372f27e2be1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/RedisTypeSelector.test.tsx @@ -1,12 +1,10 @@ import { describe, expect, it } from "vitest"; import RedisTypeSelector from "./RedisTypeSelector"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; describe("RedisTypeSelector", () => { it("should render the component", () => { - const { getAllByText } = render( - {}} />, - ); - expect(getAllByText(/Redis/i).length).toBeGreaterThan(0); + render( {}} />); + expect(screen.getAllByText(/Redis/i).length).toBeGreaterThan(0); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx index 76287e6e724..563c684237a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/CoordinationRedisTypeSelector.test.tsx @@ -5,6 +5,7 @@ import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import CoordinationRedisTypeSelector from "./CoordinationRedisTypeSelector"; import { COORDINATION_REDIS_TYPE_DESCRIPTIONS } from "./coordinationRedisFields"; +import { chooseSelectOption } from "../../../../../../tests/test-utils"; describe("CoordinationRedisTypeSelector", () => { it("labels the control and shows the current selection", () => { @@ -36,8 +37,7 @@ describe("CoordinationRedisTypeSelector", () => { const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Cluster")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Cluster"); expect(onTypeChange).toHaveBeenCalledTimes(1); expect(onTypeChange.mock.calls[0][0]).toBe("cluster"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index 8c36b934789..f320d8e0f97 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { DailyData, KeyMetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; @@ -79,25 +79,25 @@ const renderWith = (results: DailyData[], overrides: Partial describe("CacheLeakageCard", () => { it("ranks leaking keys by uncached prompt tokens and shows cache hit ratio", () => { - const { getByText, getByLabelText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-caching": key("caching-key", { prompt_tokens: 1000, cache_read_input_tokens: 900 }), "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }), ]); - expect(getByText("leaky-key")).toBeInTheDocument(); - expect(getByText("0.0%")).toBeInTheDocument(); - expect(getByText("90.0%")).toBeInTheDocument(); + expect(screen.getByText("leaky-key")).toBeInTheDocument(); + expect(screen.getByText("0.0%")).toBeInTheDocument(); + expect(screen.getByText("90.0%")).toBeInTheDocument(); [ "Input tokens you sent in this range that weren't served from or written to the cache", "Share of your input tokens that were served from the cache", "About how much you'd save if this uncached input used prompt caching. Estimated as uncached input tokens times what your cached traffic already nets per cached token (realized cache savings, after write premiums, ÷ cache read and write tokens). Blank when caching is not currently saving anything overall.", - ].forEach((info) => expect(getByLabelText(info)).toBeInTheDocument()); + ].forEach((info) => expect(screen.getByLabelText(info)).toBeInTheDocument()); }); it("sorts by the clicked column, worst cache hit rate first", () => { - const { getAllByRole, getByText } = renderWith([ + renderWith([ dayWithKeys("2026-07-12", { "hash-a": key("alpha", { prompt_tokens: 10000, @@ -111,48 +111,48 @@ describe("CacheLeakageCard", () => { }), }), ]); - const firstDataRow = () => getAllByRole("row")[1]; + const firstDataRow = () => screen.getAllByRole("row")[1]; expect(firstDataRow()).toHaveTextContent("alpha"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("bravo"); - fireEvent.click(getByText("Cache hit rate")); + fireEvent.click(screen.getByText("Cache hit rate")); expect(firstDataRow()).toHaveTextContent("alpha"); }); it("switches to the model view and lists only Anthropic models", () => { - const { getByText, queryByText } = renderWith([ + renderWith([ dayWithModels("2026-07-12", { "claude-sonnet-5": { prompt_tokens: 5000, cache_read_input_tokens: 0 }, "gpt-4o": { prompt_tokens: 8000, cache_read_input_tokens: 0 }, }), ]); - fireEvent.click(getByText("By model")); + fireEvent.click(screen.getByText("By model")); - expect(getByText("Cache leakage by model")).toBeInTheDocument(); - expect(getByText("claude-sonnet-5")).toBeInTheDocument(); - expect(queryByText("gpt-4o")).not.toBeInTheDocument(); + expect(screen.getByText("Cache leakage by model")).toBeInTheDocument(); + expect(screen.getByText("claude-sonnet-5")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o")).not.toBeInTheDocument(); }); it("shows an empty state when no key used tokens in the range", () => { - const { getByText, queryByRole } = renderWith([dayWithKeys("2026-07-12", {})]); + renderWith([dayWithKeys("2026-07-12", {})]); - expect(getByText("No key usage in this range.")).toBeInTheDocument(); - expect(queryByRole("table")).not.toBeInTheDocument(); + expect(screen.getByText("No key usage in this range.")).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); }); it("tells the user the table is still filling in while fallback pages stream", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { getByText, getByRole } = renderWith([day], { isFetchingMore: true }); + renderWith([day], { isFetchingMore: true }); - expect(getByRole("table")).toBeInTheDocument(); + expect(screen.getByRole("table")).toBeInTheDocument(); expect( - getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.getByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).toBeInTheDocument(); }); @@ -160,10 +160,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day], { loading: true }); + renderWith([day], { loading: true }); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); @@ -171,10 +171,10 @@ describe("CacheLeakageCard", () => { const day = dayWithKeys("2026-07-12", { "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), }); - const { queryByText } = renderWith([day]); + renderWith([day]); expect( - queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), + screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index 1cc7bec13d1..03250e3e53b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, waitFor } from "@testing-library/react"; +import { fireEvent, render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -56,7 +56,7 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { getByRole, getByTestId, findByTestId, queryByText } = render( + render( , @@ -64,12 +64,12 @@ describe("CostOptimizationView daily activity", () => { await waitFor(() => expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1)); - fireEvent.click(getByRole("tab", { name: "Prompt Caching" })); - await findByTestId("caching-settings"); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Caching" })); + await screen.findByTestId("caching-settings"); expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledTimes(1); expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); - expect(queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Currently fetching spend data/)).not.toBeInTheDocument(); }); it("shows the fetch-progress banner while the paginated fallback streams pages in", async () => { @@ -84,13 +84,13 @@ describe("CostOptimizationView daily activity", () => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); - const { findByText, getByRole } = render( + render( , ); - expect(await findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); - expect(getByRole("button", { name: "Stop" })).toBeInTheDocument(); + expect(await screen.findByText(/Currently fetching spend data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Stop" })).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index d5df5aa75da..028367555a1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -45,32 +45,32 @@ describe("CostOptimizationView", () => { }); it("renders the standard page header with the sidebar's Cost Optimization icon", () => { - const { container, getByRole, getByText } = renderView(); + const { container } = renderView(); - expect(getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); - expect(getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); + expect(screen.getByRole("heading", { level: 1, name: "Cost Optimization" })).toBeInTheDocument(); + expect(screen.getByText(/Track and configure the mechanisms that save you money/)).toBeInTheDocument(); expect(container.querySelector(".lucide-piggy-bank")).not.toBeNull(); }); it("renders the four cost-optimization tabs", () => { - const { getByText } = renderView(); + renderView(); - expect(getByText("Overall")).toBeInTheDocument(); - expect(getByText("Prompt Compression")).toBeInTheDocument(); - expect(getByText("Prompt Caching")).toBeInTheDocument(); - expect(getByText("Auto-Router")).toBeInTheDocument(); + expect(screen.getByText("Overall")).toBeInTheDocument(); + expect(screen.getByText("Prompt Compression")).toBeInTheDocument(); + expect(screen.getByText("Prompt Caching")).toBeInTheDocument(); + expect(screen.getByText("Auto-Router")).toBeInTheDocument(); }); it("defaults to the Overall tab and switches the active tab on click", () => { - const { getByRole } = renderView(); + renderView(); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); - fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); + fireEvent.click(screen.getByRole("tab", { name: "Prompt Compression" })); - expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); - expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); // Unlike the other three pages in this cleanup, Cost Optimization keeps its @@ -80,21 +80,21 @@ describe("CostOptimizationView", () => { // are proxy-admin-only, so those are what disappear. describe("proxy-admin-only tabs", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])("shows %s the Overall tab only", (userRole) => { - const { getByRole, queryByRole } = renderView(userRole); + renderView(userRole); - expect(getByRole("tab", { name: "Overall" })).toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Overall" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); }); it("never mounts the panels behind the admin-only endpoints for an internal user", () => { - const { getByTestId, queryByTestId } = renderView("Internal User"); + renderView("Internal User"); - expect(getByTestId("usage-tab")).toBeInTheDocument(); - expect(queryByTestId("compression-tab")).not.toBeInTheDocument(); - expect(queryByTestId("caching-tab")).not.toBeInTheDocument(); - expect(queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); + expect(screen.getByTestId("usage-tab")).toBeInTheDocument(); + expect(screen.queryByTestId("compression-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("caching-tab")).not.toBeInTheDocument(); + expect(screen.queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx index 38517dab0ab..2c602033171 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.test.tsx @@ -1,4 +1,4 @@ -import { render, waitFor } from "@testing-library/react"; +import { render, waitFor, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; const mockGetGeneralSettingsCall = vi.fn(); @@ -37,10 +37,10 @@ describe("PromptCachingTab", () => { cancelled: false, cancel: vi.fn(), }; - const { getByTestId } = render(); + render(); - expect(getByTestId("caching-settings")).toBeInTheDocument(); - expect(getByTestId("cache-leakage-card")).toBeInTheDocument(); + expect(screen.getByTestId("caching-settings")).toBeInTheDocument(); + expect(screen.getByTestId("cache-leakage-card")).toBeInTheDocument(); await waitFor(() => expect(mockCacheLeakageCard).toHaveBeenCalledWith(expect.objectContaining({ activity }))); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index 8b397f20552..a1de608d0bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -97,11 +97,13 @@ import { useStopShadowEval, type ShadowEvalJob, } from "./useShadowEval"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const job = (overrides: Partial = {}): ShadowEvalJob => ({ job_id: "job-1", status: "running", router_name: "claude-auto", + router_names: ["claude-auto"], direction: "forward", baseline_model: null, judge_model: "anthropic/claude-sonnet-5", @@ -436,8 +438,7 @@ describe("ShadowEvalSection", () => { await user.click(within(keyList).getByText("prod-alpha")); await user.click(keyInput); await user.click(within(keyList).getByText("staging-beta")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); expect(screen.getByText("Start shadow eval")).toBeDisabled(); @@ -449,7 +450,7 @@ describe("ShadowEvalSection", () => { api_key_ids: ["hash-alpha", "hash-beta"], team_ids: [], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "forward", shadow_percentage: 10, duration_days: 7, @@ -469,8 +470,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search teams by alias")); const teamList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(teamList).getByText("engineering")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); await user.click(screen.getByText("Start shadow eval")); @@ -479,7 +479,7 @@ describe("ShadowEvalSection", () => { api_key_ids: [], team_ids: ["team-eng"], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "forward", shadow_percentage: 10, duration_days: 7, @@ -501,8 +501,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByPlaceholderText("Search keys by alias")); const keyList = await screen.findByTestId("paginated-multi-select-list"); await user.click(within(keyList).getByText("prod-alpha")); - await user.click(screen.getByPlaceholderText("Select an auto-router")); - await user.click(await screen.findByText("gpt-auto")); + await chooseSelectOption(user, screen.getByPlaceholderText("Select up to 4 auto-routers"), "gpt-auto"); await user.click(screen.getByPlaceholderText("Select a judge model")); await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); @@ -517,7 +516,7 @@ describe("ShadowEvalSection", () => { api_key_ids: ["hash-alpha"], team_ids: [], user_ids: [], - router_name: "gpt-auto", + router_names: ["gpt-auto"], direction: "reverse", baseline_model: "prod-claude", shadow_percentage: 10, @@ -528,6 +527,119 @@ describe("ShadowEvalSection", () => { expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); + it("submits every picked auto-router so one job compares them on the same traffic", async () => { + const user = userEvent.setup(); + const { start } = mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Search keys by alias")); + const keyList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(keyList).getByText("prod-alpha")); + const routerInput = screen.getByPlaceholderText("Select up to 4 auto-routers"); + await user.click(routerInput); + await user.click(await screen.findByText("gpt-auto")); + await user.click(routerInput); + await user.click(await screen.findByText("claude-auto")); + expect( + screen.getByText("Every router sees the same sampled requests, judged against the same live responses"), + ).toBeInTheDocument(); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(screen.getByText("Start shadow eval")); + + const expectedBody = { + api_key_ids: ["hash-alpha"], + team_ids: [], + user_ids: [], + router_names: ["gpt-auto", "claude-auto"], + direction: "forward", + shadow_percentage: 10, + duration_days: 7, + max_budget: 10, + judge_model: "anthropic/claude-sonnet-5", + }; + expect(start.mutate).toHaveBeenCalledWith(expectedBody); + }); + + it("blocks starting a reverse job with more than one router and says why", async () => { + const user = userEvent.setup(); + mockHooks({}); + render(); + + await user.click(screen.getByPlaceholderText("Search keys by alias")); + const keyList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(keyList).getByText("prod-alpha")); + const routerInput = screen.getByPlaceholderText("Select up to 4 auto-routers"); + await user.click(routerInput); + await user.click(await screen.findByText("gpt-auto")); + await user.click(routerInput); + await user.click(await screen.findByText("claude-auto")); + await user.click(screen.getByText("Adoption check: key's traffic vs the router")); + await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + await user.click(screen.getByPlaceholderText("Select a baseline model")); + await user.click(screen.getByRole("option", { name: /prod-claude/ })); + + expect(screen.getByText("A regression check compares one router to its baseline")).toBeInTheDocument(); + expect(screen.getByText("Start shadow eval")).toBeDisabled(); + }); + + it("renders a per-router comparison table only when the job ran several routers", () => { + const routerSlice = (group: string, wins: number) => ({ + group, + turn_count: 20, + real_win_rate_pct: 100 - wins - 10, + shadow_win_rate_pct: wins, + tie_rate_pct: 10, + avg_judge_confidence: 0.8, + real_spend: 0.4, + shadow_spend: 0.2, + cache_hit_turns: 0, + }); + const base = job(); + const multi = job({ + router_names: ["claude-auto", "gpt-auto"], + results: { ...base.results!, by_router: [routerSlice("claude-auto", 40), routerSlice("gpt-auto", 70)] }, + }); + mockHooks({ jobs: [multi], detailsById: { "job-1": multi } }); + render(); + + expect(screen.getByText("Router")).toBeInTheDocument(); + const rows = screen.getAllByRole("row").map((row) => row.textContent ?? ""); + expect(rows.some((text) => text.includes("claude-auto") && text.includes("40.0%"))).toBe(true); + expect(rows.some((text) => text.includes("gpt-auto") && text.includes("70.0%"))).toBe(true); + expect( + screen.getByText( + (_, element) => + element?.textContent === "Shadowing 10% of prod-alpha traffic via claude-auto, gpt-auto" && + element.tagName === "P", + ), + ).toBeInTheDocument(); + }); + + it("renders a job from an older proxy that predates router_names", () => { + const legacy = { ...job(), router_names: undefined } as unknown as ShadowEvalJob; + mockHooks({ jobs: [legacy], detailsById: { "job-1": legacy } }); + render(); + + expect( + screen.getByText( + (_, element) => + element?.textContent === "Shadowing 10% of prod-alpha traffic via claude-auto" && element.tagName === "P", + ), + ).toBeInTheDocument(); + }); + + it("keeps the per-router table hidden for a single-router job", () => { + const base = job(); + const single = job({ results: { ...base.results!, by_router: [] } }); + mockHooks({ jobs: [single], detailsById: { "job-1": single } }); + render(); + + expect(screen.queryByText("Router")).not.toBeInTheDocument(); + }); + it("flips the arm labels and headline for a reverse job's results", () => { const j = job({ direction: "reverse", baseline_model: "openai/gpt-4o" }); mockHooks({ jobs: [j], detailsById: { "job-1": j } }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index ec145bc617a..c66d74074c2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -2,32 +2,21 @@ import React, { useMemo, useState } from "react"; -import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; -import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; -import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; -import TeamMultiSelect from "@/components/common_components/team_multi_select"; -import { userOptionLabel } from "@/components/common_components/UserDropdown"; -import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { CircleHelp } from "lucide-react"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card } from "@/components/ui/card"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { Input } from "@/components/ui/input"; -import { Label } from "@/components/ui/label"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { ApiError } from "@/lib/http/client"; import { usd } from "./costOptimizationUtils"; +import { StartForm } from "./ShadowEvalStartForm"; import { useShadowEvalJob, useShadowEvalJobs, - useStartShadowEval, useStopShadowEval, type ShadowEvalJob, type ShadowEvalJobTarget, @@ -96,17 +85,19 @@ const targetStatus = (job: ShadowEvalJob, target: ShadowEvalJobTarget): string = return target.stopped_at != null ? "stopped" : "running"; }; +const jobRouters = (job: ShadowEvalJob): string => (job.router_names ?? [job.router_name]).join(", "); + const jobHeadline = (job: ShadowEvalJob): React.ReactNode => job.direction === "reverse" ? ( <> - Comparing {job.router_name} to{" "} + Comparing {jobRouters(job)} to{" "} {job.baseline_model} on {job.shadow_percentage}% of{" "} {shadowedTargetsLabel(job)} traffic ) : ( <> Shadowing {job.shadow_percentage}% of {shadowedTargetsLabel(job)}{" "} - traffic via {job.router_name} + traffic via {jobRouters(job)} ); @@ -352,6 +343,11 @@ const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({
+ {(results.by_router ?? []).length > 1 && ( +
+ +
+ )} {results.by_current_model.length > 0 && ( { - const { data: costMap } = useModelCostMap(); - return useMemo(() => { - if (!costMap) return []; - const chatModels = Object.entries(costMap as Record) - .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) - .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); - return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); - }, [costMap]); -}; - -const useJudgeModelOptions = (): SearchSelectOption[] => { - const chatModels = useChatModelNames(); - return useMemo(() => { - const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ - label: model, - value: model, - sublabel: "Recommended", - })); - const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); - const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); - return [...pinned, ...rest]; - }, [chatModels]); -}; - -const useBaselineModelOptions = (): SearchSelectOption[] => { - const configuredGroups = usePlainModelGroups(); - const chatModels = useChatModelNames(); - return useMemo(() => { - const configured = [...configuredGroups] - .toSorted((a, b) => a.localeCompare(b)) - .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); - const rest = chatModels - .filter((model) => !configuredGroups.has(model)) - .map((model) => ({ label: model, value: model })); - return [...configured, ...rest]; - }, [configuredGroups, chatModels]); -}; - -const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ - { value: "forward", label: "Adoption check: key's traffic vs the router" }, - { value: "reverse", label: "Regression check: router's picks vs a baseline" }, -] as const; - -const START_FORM_DESCRIPTION: Record = { - forward: - "Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.", - reverse: - "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.", -}; - -const DURATION_OPTIONS = [ - { value: "1", label: "1 day" }, - { value: "3", label: "3 days" }, - { value: "7", label: "7 days" }, - { value: "14", label: "14 days" }, - { value: "30", label: "30 days" }, -] as const; - -const Field: React.FC<{ label: string; htmlFor?: string; className?: string; children: React.ReactNode }> = ({ - label, - htmlFor, - className, - children, -}) => ( -
- - {children} -
-); - -const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => { - const [search, setSearch] = useState(""); - const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, { - selectedKeyAlias: search || null, - }); - const options = useMemo( - () => - (data?.pages ?? []) - .flatMap((page) => page.keys) - .map((key) => ({ - label: key.key_alias || key.key_name || key.token, - value: key.token, - sublabel: key.token, - })), - [data], - ); - return ( - void fetchNextPage()} - hasNextPage={hasNextPage} - isFetchingNextPage={isFetchingNextPage} - isLoading={isPending} - placeholder="Search keys by alias" - emptyText="No matching keys" - errorText={isError ? "Keys could not be loaded. Refresh the page to retry." : undefined} - /> - ); -}; - -const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => { - const [search, setSearch] = useState(""); - const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers( - 50, - search || undefined, - ); - const options = useMemo( - () => - Array.from( - new Map( - (data?.pages ?? []) - .flatMap((page) => page.users) - .map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const), - ).values(), - ), - [data], - ); - return ( - void fetchNextPage()} - hasNextPage={hasNextPage} - isFetchingNextPage={isFetchingNextPage} - isLoading={isPending} - placeholder="Search users by email" - emptyText="No matching users" - errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined} - /> - ); -}; - -const StartForm: React.FC = () => { - const { accessToken } = useAuthorized(); - const [apiKeyIds, setApiKeyIds] = useState([]); - const [teamIds, setTeamIds] = useState([]); - const [userIds, setUserIds] = useState([]); - const [routerName, setRouterName] = useState(""); - const [direction, setDirection] = useState("forward"); - const [baselineModel, setBaselineModel] = useState(""); - const [percentage, setPercentage] = useState("10"); - const [durationDays, setDurationDays] = useState("7"); - const [judgeModel, setJudgeModel] = useState(""); - const [maxBudget, setMaxBudget] = useState("10"); - const { data: autoRouters } = useAutoRouters(); - const judgeModelOptions = useJudgeModelOptions(); - const baselineModelOptions = useBaselineModelOptions(); - const start = useStartShadowEval(); - - const routerOptions = useMemo(() => { - const names = new Set( - (autoRouters ?? []).map((deployment) => deployment.model_name).filter((name): name is string => Boolean(name)), - ); - return [...names].toSorted().map((name) => ({ label: name, value: name })); - }, [autoRouters]); - - const parsedPct = Number.parseFloat(percentage); - const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; - const parsedMaxBudget = Number.parseFloat(maxBudget); - const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; - const baselinePicked = direction === "forward" || baselineModel !== ""; - const targetsPicked = apiKeyIds.length + teamIds.length + userIds.length > 0; - const filled = targetsPicked && [routerName, judgeModel].every((field) => field !== "") && baselinePicked; - const boundsValid = percentageValid && maxBudgetValid; - const valid = Boolean(accessToken) && filled && boundsValid; - const handleStart = () => { - const startBody = { - api_key_ids: apiKeyIds, - team_ids: teamIds, - user_ids: userIds, - router_name: routerName, - direction, - ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), - shadow_percentage: parsedPct, - duration_days: Number.parseInt(durationDays, 10), - max_budget: parsedMaxBudget, - judge_model: judgeModel, - }; - start.mutate(startBody); - }; - - return ( - - - Start a shadow eval -

{START_FORM_DESCRIPTION[direction]}

- - -
- - - - - - - - - - - - - - - - -
- setPercentage(e.target.value)} - /> - % of traffic -
-
- {percentage.trim() !== "" && !percentageValid && ( -

Enter a value from 0.1 to 100

- )} -
- - - - - -
- $ - setMaxBudget(e.target.value)} - /> - max shadow + judge spend, per target -
- {maxBudget.trim() !== "" && !maxBudgetValid && ( -

Enter a value from 0.01 to 10000

- )} - - {direction === "reverse" && ( - - - - )} - - - -
- - - - ); -}; - const previousSummary = (job: ShadowEvalJob): string => { const results = job.results; if (results) return pct(routerMatchedOrBeatPct(job.direction, results)); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx new file mode 100644 index 00000000000..f96910a4ad6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalStartForm.tsx @@ -0,0 +1,432 @@ +"use client"; + +import React, { useMemo, useState } from "react"; + +import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; +import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; +import TeamMultiSelect from "@/components/common_components/team_multi_select"; +import { userOptionLabel } from "@/components/common_components/UserDropdown"; +import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; + +import { useStartShadowEval, type ShadowEvalJob } from "./useShadowEval"; + +type ShadowEvalDirection = ShadowEvalJob["direction"]; + +const MAX_ROUTERS = 4; + +const RECOMMENDED_JUDGE_MODELS = ["anthropic/claude-sonnet-5", "openai/gpt-4o", "gemini/gemini-2.5-pro"] as const; + +interface CostMapEntry { + litellm_provider?: string; + mode?: string; +} + +const useChatModelNames = (): string[] => { + const { data: costMap } = useModelCostMap(); + return useMemo(() => { + if (!costMap) return []; + const chatModels = Object.entries(costMap as Record) + .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) + .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); + return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); + }, [costMap]); +}; + +const useJudgeModelOptions = (): SearchSelectOption[] => { + const chatModels = useChatModelNames(); + return useMemo(() => { + const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ + label: model, + value: model, + sublabel: "Recommended", + })); + const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); + const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); + return [...pinned, ...rest]; + }, [chatModels]); +}; + +const useBaselineModelOptions = (): SearchSelectOption[] => { + const configuredGroups = usePlainModelGroups(); + const chatModels = useChatModelNames(); + return useMemo(() => { + const configured = [...configuredGroups] + .toSorted((a, b) => a.localeCompare(b)) + .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); + const rest = chatModels + .filter((model) => !configuredGroups.has(model)) + .map((model) => ({ label: model, value: model })); + return [...configured, ...rest]; + }, [configuredGroups, chatModels]); +}; + +const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ + { value: "forward", label: "Adoption check: key's traffic vs the router" }, + { value: "reverse", label: "Regression check: router's picks vs a baseline" }, +] as const; + +const START_FORM_DESCRIPTION: Record = { + forward: + "Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.", + reverse: + "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.", +}; + +const DURATION_OPTIONS = [ + { value: "1", label: "1 day" }, + { value: "3", label: "3 days" }, + { value: "7", label: "7 days" }, + { value: "14", label: "14 days" }, + { value: "30", label: "30 days" }, +] as const; + +const Field: React.FC<{ label: string; htmlFor?: string; className?: string; children: React.ReactNode }> = ({ + label, + htmlFor, + className, + children, +}) => ( +
+ + {children} +
+); + +const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => { + const [search, setSearch] = useState(""); + const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, { + selectedKeyAlias: search || null, + }); + const options = useMemo( + () => + (data?.pages ?? []) + .flatMap((page) => page.keys) + .map((key) => ({ + label: key.key_alias || key.key_name || key.token, + value: key.token, + sublabel: key.token, + })), + [data], + ); + return ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isFetchingNextPage={isFetchingNextPage} + isLoading={isPending} + placeholder="Search keys by alias" + emptyText="No matching keys" + errorText={isError ? "Keys could not be loaded. Refresh the page to retry." : undefined} + /> + ); +}; + +const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => { + const [search, setSearch] = useState(""); + const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers( + 50, + search || undefined, + ); + const options = useMemo( + () => + Array.from( + new Map( + (data?.pages ?? []) + .flatMap((page) => page.users) + .map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const), + ).values(), + ), + [data], + ); + return ( + void fetchNextPage()} + hasNextPage={hasNextPage} + isFetchingNextPage={isFetchingNextPage} + isLoading={isPending} + placeholder="Search users by email" + emptyText="No matching users" + errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined} + /> + ); +}; + +const RouterField: React.FC<{ + options: SearchSelectOption[]; + routerNames: string[]; + onChange: (names: string[]) => void; + direction: ShadowEvalDirection; +}> = ({ options, routerNames, onChange, direction }) => ( + + + {routerNames.length > MAX_ROUTERS && ( +

Pick at most {MAX_ROUTERS} auto-routers

+ )} + {direction === "reverse" && routerNames.length > 1 && ( +

A regression check compares one router to its baseline

+ )} + {direction === "forward" && routerNames.length > 1 && ( +

+ Every router sees the same sampled requests, judged against the same live responses +

+ )} + +); + +interface StartFormValidityInputs { + accessToken: string | null | undefined; + apiKeyIds: string[]; + teamIds: string[]; + userIds: string[]; + routerNames: string[]; + direction: ShadowEvalDirection; + baselineModel: string; + judgeModel: string; + percentage: string; + maxBudget: string; +} + +const startFormValidity = (inputs: StartFormValidityInputs) => { + const parsedPct = Number.parseFloat(inputs.percentage); + const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; + const parsedMaxBudget = Number.parseFloat(inputs.maxBudget); + const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000; + const baselinePicked = inputs.direction === "forward" || inputs.baselineModel !== ""; + const targetsPicked = inputs.apiKeyIds.length + inputs.teamIds.length + inputs.userIds.length > 0; + const routerCountValid = inputs.routerNames.length >= 1 && inputs.routerNames.length <= MAX_ROUTERS; + const routersMatchDirection = inputs.direction === "forward" || inputs.routerNames.length === 1; + const routersValid = routerCountValid && routersMatchDirection; + const modelsPicked = routersValid && inputs.judgeModel !== "" && baselinePicked; + const filled = targetsPicked && modelsPicked; + const boundsValid = percentageValid && maxBudgetValid; + const valid = Boolean(inputs.accessToken) && filled && boundsValid; + return { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid }; +}; + +interface StartBodyInputs { + apiKeyIds: string[]; + teamIds: string[]; + userIds: string[]; + routerNames: string[]; + direction: ShadowEvalDirection; + baselineModel: string; + shadowPercentage: number; + durationDays: number; + maxBudget: number; + judgeModel: string; +} + +const buildStartBody = (inputs: StartBodyInputs) => ({ + api_key_ids: inputs.apiKeyIds, + team_ids: inputs.teamIds, + user_ids: inputs.userIds, + router_names: inputs.routerNames, + direction: inputs.direction, + ...(inputs.direction === "reverse" ? { baseline_model: inputs.baselineModel } : {}), + shadow_percentage: inputs.shadowPercentage, + duration_days: inputs.durationDays, + max_budget: inputs.maxBudget, + judge_model: inputs.judgeModel, +}); + +export const StartForm: React.FC = () => { + const { accessToken } = useAuthorized(); + const [apiKeyIds, setApiKeyIds] = useState([]); + const [teamIds, setTeamIds] = useState([]); + const [userIds, setUserIds] = useState([]); + const [routerNames, setRouterNames] = useState([]); + const [direction, setDirection] = useState("forward"); + const [baselineModel, setBaselineModel] = useState(""); + const [percentage, setPercentage] = useState("10"); + const [durationDays, setDurationDays] = useState("7"); + const [judgeModel, setJudgeModel] = useState(""); + const [maxBudget, setMaxBudget] = useState("10"); + const { data: autoRouters } = useAutoRouters(); + const judgeModelOptions = useJudgeModelOptions(); + const baselineModelOptions = useBaselineModelOptions(); + const start = useStartShadowEval(); + + const routerOptions = useMemo(() => { + const names = new Set( + (autoRouters ?? []).map((deployment) => deployment.model_name).filter((name): name is string => Boolean(name)), + ); + return [...names].toSorted().map((name) => ({ label: name, value: name })); + }, [autoRouters]); + + const validityInputs: StartFormValidityInputs = { + accessToken, + apiKeyIds, + teamIds, + userIds, + routerNames, + direction, + baselineModel, + judgeModel, + percentage, + maxBudget, + }; + const { parsedPct, parsedMaxBudget, percentageValid, maxBudgetValid, valid } = startFormValidity(validityInputs); + const handleStart = () => { + const bodyInputs: StartBodyInputs = { + apiKeyIds, + teamIds, + userIds, + routerNames, + direction, + baselineModel, + shadowPercentage: parsedPct, + durationDays: Number.parseInt(durationDays, 10), + maxBudget: parsedMaxBudget, + judgeModel, + }; + start.mutate(buildStartBody(bodyInputs)); + }; + + return ( + + + Start a shadow eval +

{START_FORM_DESCRIPTION[direction]}

+ + +
+ + + + + + + + + + + + + + +
+ setPercentage(e.target.value)} + /> + % of traffic +
+
+ {percentage.trim() !== "" && !percentageValid && ( +

Enter a value from 0.1 to 100

+ )} +
+ + + + + +
+ $ + setMaxBudget(e.target.value)} + /> + max shadow + judge spend, per target +
+ {maxBudget.trim() !== "" && !maxBudgetValid && ( +

Enter a value from 0.01 to 10000

+ )} + + {direction === "reverse" && ( + + + + )} + + + +
+ + + + ); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index df23e5509bf..f85a667a074 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ToolSpendResponse } from "@/components/networking"; @@ -152,13 +152,13 @@ describe("UsageTab", () => { gateway_injected_caching_savings_spend: 0.006, compression_saved_tokens: 100000, }; - const { getByText } = renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); + renderWith([day("2026-07-12", firstDay), day("2026-07-13", secondDay)]); - expect(getByText("$0.1500")).toBeInTheDocument(); - expect(getByText("$0.1400")).toBeInTheDocument(); - expect(getByText("$0.0100")).toBeInTheDocument(); - expect(getByText("$0.0160")).toBeInTheDocument(); - expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + expect(screen.getByText("$0.1500")).toBeInTheDocument(); + expect(screen.getByText("$0.1400")).toBeInTheDocument(); + expect(screen.getByText("$0.0100")).toBeInTheDocument(); + expect(screen.getByText("$0.0160")).toBeInTheDocument(); + expect(screen.getByText("140,000 tokens compressed")).toBeInTheDocument(); }); const twoDays = () => [ @@ -167,11 +167,11 @@ describe("UsageTab", () => { ]; it("opens on a running total anchored at $0 at the start of the range", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative prepends a synthetic $0 point at the range start (Jul 1) so the // line rises from zero rather than floating; the daily running totals follow. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(3); expect(series[0]).toMatchObject({ date: "Jul 1", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); @@ -183,12 +183,12 @@ describe("UsageTab", () => { // The original complaint: a one-day range plotted a single floating dot. The // synthetic start anchor gives the line a zero origin to climb from. const oneDay = new Date(2026, 6, 24); - const { getByTestId } = renderWith( - [day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], - { from: oneDay, to: oneDay }, - ); + renderWith([day("2026-07-24", { compression_savings_spend: 0.2, gateway_injected_caching_savings_spend: 0.05 })], { + from: oneDay, + to: oneDay, + }); - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ date: "Jul 24", Compression: 0, "Prompt caching": 0 }); expect(series[1]).toMatchObject({ date: "Jul 24", Compression: 0.2, "Prompt caching": 0.05 }); @@ -202,49 +202,49 @@ describe("UsageTab", () => { day("2026-07-13", { gateway_injected_caching_savings_spend: 0.1 }), day("2026-07-12", { gateway_injected_caching_savings_spend: 0.04 }), ]; - const { getByTestId, getByRole } = renderWith(newestFirst); + renderWith(newestFirst); // The $0 anchor leads, then the days climb oldest to newest. - const cumulative = readSeries(getByTestId("area-chart")); + const cumulative = readSeries(screen.getByTestId("area-chart")); expect(cumulative.map((p: { date: string }) => p.date)).toEqual(["Jul 1", "Jul 12", "Jul 13"]); expect(cumulative[1]["Prompt caching"]).toBeCloseTo(0.04, 5); expect(cumulative[2]["Prompt caching"]).toBeCloseTo(0.14, 5); expect(cumulative[2]["Prompt caching"]).toBeGreaterThan(cumulative[1]["Prompt caching"]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const perDay = readSeries(getByTestId("bar-chart")); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const perDay = readSeries(screen.getByTestId("bar-chart")); expect(perDay.map((p: { date: string }) => p.date)).toEqual(["Jul 12", "Jul 13"]); }); it("draws bars of the raw per-interval readings on the other tab", async () => { - const { getByRole, getByTestId, queryByTestId } = renderWith(twoDays()); + renderWith(twoDays()); // Cumulative opens on the area line. - expect(getByTestId("area-chart")).toBeInTheDocument(); + expect(screen.getByTestId("area-chart")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); // Per day switches to a bar chart of the unaccumulated daily savings, with no // synthetic anchor prepended. - expect(queryByTestId("area-chart")).not.toBeInTheDocument(); - const series = readSeries(getByTestId("bar-chart")); + expect(screen.queryByTestId("area-chart")).not.toBeInTheDocument(); + const series = readSeries(screen.getByTestId("bar-chart")); expect(series).toHaveLength(2); expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); }); it("says what the line means and over what range", async () => { - const { getByText, getByRole } = renderWith(twoDays()); + renderWith(twoDays()); - expect(getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); - await userEvent.click(getByRole("tab", { name: "Per day" })); - expect(getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + expect(screen.getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + expect(screen.getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); }); it("builds the per-driver donut from the range totals, not the running total", () => { - const { getByTestId } = renderWith(twoDays()); + renderWith(twoDays()); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -252,9 +252,9 @@ describe("UsageTab", () => { }); it("omits a driver slice when that driver has no savings", () => { - const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([{ driver: "Compression", color: "emerald", usd: expect.closeTo(0.04, 5) }]); }); @@ -262,7 +262,7 @@ describe("UsageTab", () => { // Stacking sums the series into one bar. Auto-router savings go negative when a // model switch pays for a cold cache, and that segment would be drawn below the // axis while the rest of the bar still read as the day's total. - const { getByRole, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -270,8 +270,8 @@ describe("UsageTab", () => { }), ]); - await userEvent.click(getByRole("tab", { name: "Per day" })); - const bars = getByTestId("bar-chart"); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); + const bars = screen.getByTestId("bar-chart"); expect(bars).toHaveAttribute("data-stack", "false"); expect(readSeries(bars)[0]).toMatchObject({ "Auto-router": -0.05 }); }); @@ -281,10 +281,10 @@ describe("UsageTab", () => { // per day"). Hand-rolled rows made it compete with the legend and the toggle for // width, so the header grew a line on one tab and the chart moved with it. CardHeader // sizes the action column to its content and gives the rest to the title column. - const { getByRole, getByTestId, container } = renderWith(twoDays()); + const { container } = renderWith(twoDays()); const header = () => { - const legend = getByTestId("chart-legend"); + const legend = screen.getByTestId("chart-legend"); const action = legend.closest('[data-slot="card-action"]') as HTMLElement; const cardHeader = action.parentElement as HTMLElement; const description = cardHeader.querySelector('[data-slot="card-description"]') as HTMLElement; @@ -295,12 +295,12 @@ describe("UsageTab", () => { expect(before.action).toBeTruthy(); expect(before.description).toBeTruthy(); // the toggle rides in the same action slot as the legend, so neither moves alone - expect(before.action.contains(getByRole("tablist"))).toBe(true); + expect(before.action.contains(screen.getByRole("tablist"))).toBe(true); // the subtitle lives outside that slot, so its length cannot reposition the controls expect(before.action.contains(before.description)).toBe(false); expect(before.description).toHaveTextContent(/Running total saved/); - await userEvent.click(getByRole("tab", { name: "Per day" })); + await userEvent.click(screen.getByRole("tab", { name: "Per day" })); const after = header(); expect(after.action).toBe(before.action); @@ -314,7 +314,7 @@ describe("UsageTab", () => { // Switching models leaves the new one with a cold cache, so a route can cost more // than the baseline would have. A negative slice is meaningless in a donut, but the // total has to keep the loss or the page can only ever report good news. - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.1, gateway_injected_caching_savings_spend: 0.02, @@ -322,16 +322,16 @@ describe("UsageTab", () => { }), ]); - expect(getByText("$0.0700")).toBeInTheDocument(); - expect(getByText("-$0.0500")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("-$0.0500")).toBeInTheDocument(); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices.map((d: { driver: string }) => d.driver)).toEqual(["Compression", "Prompt caching"]); - expect(getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); + expect(screen.getByTestId("donut-chart")).toHaveAttribute("data-label", "$0.1200"); }); it("carries auto-router savings into the summary card, donut slice, and cumulative series", () => { - const { getByText, getByTestId } = renderWith([ + renderWith([ day("2026-07-12", { compression_savings_spend: 0.04, gateway_injected_caching_savings_spend: 0.006, @@ -345,11 +345,11 @@ describe("UsageTab", () => { ]); // Total saved now sums three drivers, and the auto-router card carries its own total. - expect(getByText("$0.2260")).toBeInTheDocument(); - expect(getByText("$0.0700")).toBeInTheDocument(); + expect(screen.getByText("$0.2260")).toBeInTheDocument(); + expect(screen.getByText("$0.0700")).toBeInTheDocument(); // The driver donut gains a third slice priced from the range totals. - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + const slices = JSON.parse(screen.getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); expect(slices).toEqual([ { driver: "Compression", color: "emerald", usd: expect.closeTo(0.14, 5) }, { driver: "Prompt caching", color: "blue", usd: expect.closeTo(0.016, 5) }, @@ -357,7 +357,7 @@ describe("UsageTab", () => { ]); // And the cumulative line accumulates the auto-router series alongside the others. - const series = readSeries(getByTestId("area-chart")); + const series = readSeries(screen.getByTestId("area-chart")); expect(series[2]["Auto-router"]).toBeCloseTo(0.07, 5); }); @@ -371,9 +371,9 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]"); expect(series[0]).toMatchObject({ tool_name: "search", spend: 4.0 }); // The 64px bar cap is this card's opt-in; the shared BarChart must not cap @@ -391,14 +391,16 @@ describe("UsageTab", () => { start_date: "2026-07-12", end_date: "2026-07-12", }; - const { findAllByTestId, getAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend }); + renderWith([day("2026-07-12", {})], { toolSpend }); - const bars = await findAllByTestId("bar-chart"); + const bars = await screen.findAllByTestId("bar-chart"); const [totalByTool, dailyByTool] = bars.slice(-2); expect(dailyByTool).toHaveAttribute("data-show-legend", "false"); expect(totalByTool).toHaveAttribute("data-colors", dailyByTool.getAttribute("data-colors")); - const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); + const toolLegends = screen + .getAllByTestId("chart-legend") + .filter((legend) => legend.textContent === "search,read_file"); expect(toolLegends).toHaveLength(1); }); @@ -415,23 +417,23 @@ describe("UsageTab", () => { it.each(["Internal User", "Internal Viewer", "Org Admin"])( "hides the card and never calls the endpoint for %s", async (userRole) => { - const { queryByText, getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend, userRole, }); // Liveness gate: the daily-activity charts still render for this role, // so the absence below is the gate, not an empty tab. - expect(getByTestId("donut-chart")).toBeInTheDocument(); - expect(queryByText("Spend by tool")).not.toBeInTheDocument(); + expect(screen.getByTestId("donut-chart")).toBeInTheDocument(); + expect(screen.queryByText("Spend by tool")).not.toBeInTheDocument(); await vi.waitFor(() => expect(mockGetToolSpend).not.toHaveBeenCalled()); }, ); it("keeps the card and the endpoint call for an admin", async () => { - const { findByText } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); + renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); - expect(await findByText("Spend by tool")).toBeInTheDocument(); + expect(await screen.findByText("Spend by tool")).toBeInTheDocument(); expect(mockGetToolSpend).toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 03cfeed42ff..7785a8e44ab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -312,4 +312,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + alice: { + provider: "Alice", + guardrailNameSuggestion: "Alice", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts index 13909e48185..1e486639840 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.test.ts @@ -27,6 +27,7 @@ const EXPECTED_PARTNER_LOGO_FILES: Record = { deepkeep: "deepkeep.svg", repelloai: "repelloai.png", straiker: "straiker.svg", + alice: "alice.svg", }; describe("guardrail_garden_data logos", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index 744af89a357..931b3a111d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -464,6 +464,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Agentic", "Prompt Injection", "Tool Misuse", "MCP", "Skills"], providerKey: "Straiker", }, + { + id: "alice", + name: "Alice", + description: + "Policy-based guardrails for prompts and model responses, evaluated per application so one proxy can enforce a different policy set per team or product.", + category: "partner", + logo: guardrailLogoMap["Alice"], + tags: ["Content Moderation", "Prompt Injection", "PII", "Policy"], + providerKey: "Alice", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx index 0a0f70b6bc2..dda56a06ab1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx @@ -1,6 +1,7 @@ import React, { useState } from "react"; import { ArrowLeft } from "lucide-react"; import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/cva.config"; import AddGuardrailForm from "./add_guardrail_form"; import { Logo } from "@/components/molecules/logo/Logo"; import { GUARDRAIL_PRESETS } from "./guardrail_garden_configs"; @@ -40,7 +41,7 @@ const GuardrailDetailView: React.FC = ({ card, onBack, const tabs = [{ key: "overview", label: "Overview" }, ...(card.eval ? [{ key: "eval", label: "Eval Results" }] : [])]; return ( -
+
{/* Back link */}
= ({ card, onBack,
{/* ── Header block (Vertex-style) ── */} -
+
-

{card.name}

+

{card.name}

-

{card.description}

+

{card.description}

{/* Action buttons — outlined style like Vertex */}
@@ -66,21 +67,18 @@ const GuardrailDetailView: React.FC = ({ card, onBack,
{/* ── Tab bar ──────────────────────────────────── */} -
-
+
+
{tabs.map((tab) => (
setActiveTab(tab.key)} - style={{ - padding: "12px 20px", - fontSize: 14, - color: activeTab === tab.key ? "#1a73e8" : "#5f6368", - borderBottom: activeTab === tab.key ? "3px solid #1a73e8" : "3px solid transparent", - cursor: "pointer", - fontWeight: activeTab === tab.key ? 500 : 400, - marginBottom: -1, - }} + className={cn( + "-mb-px cursor-pointer border-b-[3px] px-5 py-3 text-sm", + activeTab === tab.key + ? "border-info font-medium text-info" + : "border-transparent font-normal text-muted-foreground", + )} > {tab.label}
@@ -90,31 +88,27 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* ── Tab content ──────────────────────────────── */} {activeTab === "overview" && ( -
+
{/* Left column — overview + details table */} -
-

Overview

-

{card.description}

+
+

Overview

+

{card.description}

-

Guardrail Details

-

Details are as follows

+

Guardrail Details

+

Details are as follows

- +
- - - + + + {detailRows.map((row, i) => ( - - - + + + ))} @@ -122,37 +116,30 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {/* Right column — metadata sidebar like Vertex */} -
+
{/* Guardrail ID */} -
-
Guardrail ID
-
litellm/{card.id}
+
+
Guardrail ID
+
litellm/{card.id}
{/* Type */} -
-
Type
-
+
+
Type
+
{card.category === "litellm" ? "Content Filter" : "Partner"}
{/* Tags — pill style like Vertex */} {card.tags.length > 0 && ( -
-
Tags
-
+
+
Tags
+
{card.tags.map((tag) => ( {tag} @@ -166,19 +153,19 @@ const GuardrailDetailView: React.FC = ({ card, onBack, {activeTab === "eval" && (
-

Eval Results

-
- Property - - {card.name} -
Property{card.name}
{row.property}{row.value}
{row.property}{row.value}
+

Eval Results

+
- - - + + + {evalRows.map((row, i) => ( - - - + + + ))} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 2b90a1d8cbc..fcffc2122e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -1,5 +1,5 @@ import * as networking from "@/components/networking"; -import { fireEvent, render, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, waitFor, within, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import GuardrailInfoView from "./guardrail_info"; @@ -65,21 +65,19 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getAllByText, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); // Wait for the loading to complete and data to be rendered await waitFor(() => { // The guardrail name appears in multiple places (title and settings tab) - const elements = getAllByText("Test Guardrail"); + const elements = screen.getAllByText("Test Guardrail"); expect(elements.length).toBeGreaterThan(0); }); // Verify other key elements are present - expect(getByText("Back to Guardrails")).toBeInTheDocument(); - expect(getByText("Overview")).toBeInTheDocument(); - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Back to Guardrails")).toBeInTheDocument(); + expect(screen.getByText("Overview")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); it("should render a tag-based mode object rather than crashing the detail view", async () => { @@ -105,11 +103,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findAllByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); it("should render the provider logo from the bundled guardrail logo map", async () => { @@ -135,11 +131,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByAltText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - const logo = await findByAltText("Presidio PII logo"); + const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); }); @@ -167,25 +161,27 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText, findByText, container } = render( + const { container } = render( {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Click the Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); // Wait for the Settings panel to render await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); await userEvent.hover(within(container).getByRole("img", { name: "Config guardrail details" })); - expect(await findByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); + expect( + await screen.findByText("Guardrail is defined in the config file and cannot be edited."), + ).toBeInTheDocument(); }); it("should render the guardrail info", async () => { @@ -216,12 +212,10 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("PII Entity Configuration")).toBeInTheDocument(); + expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); }); }); it("should handle content filter updates correctly", async () => { @@ -251,30 +245,28 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - const { getByText, getByLabelText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); await waitFor(() => { - expect(getByText("Settings")).toBeInTheDocument(); + expect(screen.getByText("Settings")).toBeInTheDocument(); }); // Go to Settings tab - fireEvent.click(getByText("Settings")); + fireEvent.click(screen.getByText("Settings")); await waitFor(() => { - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); // Enter Edit Mode - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Modify Guardrail Name to force an update - const nameInput = getByLabelText("Guardrail Name"); + const nameInput = screen.getByLabelText("Guardrail Name"); fireEvent.change(nameInput, { target: { value: "Updated Name" } }); // Save with only name change - const saveButton = getByText("Save Changes"); + const saveButton = screen.getByText("Save Changes"); fireEvent.click(saveButton); await waitFor(() => { @@ -300,16 +292,16 @@ describe("Guardrail Info", () => { // Enter Edit Mode again to make changes await waitFor(() => { - expect(getByText("Edit Settings")).toBeInTheDocument(); + expect(screen.getByText("Edit Settings")).toBeInTheDocument(); }); - fireEvent.click(getByText("Edit Settings")); + fireEvent.click(screen.getByText("Edit Settings")); // Now modify the values using the mock button - const simulateChangeButton = getByText("Simulate Change"); + const simulateChangeButton = screen.getByText("Simulate Change"); fireEvent.click(simulateChangeButton); // Save again - fireEvent.click(getByText("Save Changes")); + fireEvent.click(screen.getByText("Save Changes")); await waitFor(() => { expect(networking.updateGuardrailCall).toHaveBeenCalled(); @@ -339,12 +331,10 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - const { findByRole, getByRole, getByText } = render( - {}} accessToken="123" isAdmin={true} />, - ); + render( {}} accessToken="123" isAdmin={true} />); - expect(await findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); - expect(getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); - expect(getByText("Guardrail Settings")).toBeInTheDocument(); + expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); + expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 83038b8e0e7..c1f2ddcf51c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -1,5 +1,6 @@ import aimSecurityLogo from "../../../../../public/assets/logos/aim_security.jpeg"; import aktoLogo from "../../../../../public/assets/logos/akto.svg"; +import aliceLogo from "../../../../../public/assets/logos/alice.svg"; import aporiaLogo from "../../../../../public/assets/logos/aporia.png"; import bedrockLogo from "../../../../../public/assets/logos/bedrock.svg"; import catoNetworksLogo from "../../../../../public/assets/logos/cato_networks.svg"; @@ -83,6 +84,7 @@ export const guardrail_provider_map: Record = { Deepkeep: "deepkeep", QostodianNexus: "qostodian_nexus", Repelloai: "repelloai", + Alice: "alice", }; // Function to populate provider map from API response - updates the original map @@ -204,6 +206,7 @@ export const guardrailLogoMap = { "Qostodian Nexus": qohashLogo.src, "RepelloAI Argus": repelloAiLogo.src, Straiker: straikerLogo.src, + Alice: aliceLogo.src, } satisfies Record; export const getGuardrailLogo = (displayName: string): string | undefined => diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx index 2f839487111..30ff2cf7c5f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import { CategoryFilter, QuickActions, PiiEntityList } from "./pii_components"; import type { PiiEntityCategory } from "@/components/guardrails/types"; @@ -6,25 +6,21 @@ import type { PiiEntityCategory } from "@/components/guardrails/types"; describe("CategoryFilter", () => { it("should render", () => { const emptyCategories: PiiEntityCategory[] = []; - const { getByText } = render( - {}} />, - ); - expect(getByText("Filter by category")).toBeInTheDocument(); + render( {}} />); + expect(screen.getByText("Filter by category")).toBeInTheDocument(); }); }); describe("QuickActions", () => { it("should render", () => { - const { getByText } = render( - {}} onUnselectAll={() => {}} hasSelectedEntities={false} />, - ); - expect(getByText("Quick Actions")).toBeInTheDocument(); + render( {}} onUnselectAll={() => {}} hasSelectedEntities={false} />); + expect(screen.getByText("Quick Actions")).toBeInTheDocument(); }); }); describe("PiiEntityList", () => { it("should render", () => { - const { getByText } = render( + render( { entityToCategoryMap={new Map()} />, ); - expect(getByText("No PII types match your filter criteria")).toBeInTheDocument(); + expect(screen.getByText("No PII types match your filter criteria")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx index 00c568ef35b..4f822578fe8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx @@ -1,10 +1,10 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, it, expect } from "vitest"; import PiiConfiguration from "./pii_configuration"; describe("PiiConfiguration", () => { it("should render", () => { - const { getByText } = render( + render( { entityCategories={[]} />, ); - expect(getByText("Configure PII Protection")).toBeInTheDocument(); + expect(screen.getByText("Configure PII Protection")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx index 26e66c8338c..81c55e7982a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolTestPanel.test.tsx @@ -6,6 +6,7 @@ import { describe, expect, it, vi } from "vitest"; import { ToolTestPanel } from "./ToolTestPanel"; import { InputSchema, MCPTool } from "@/components/mcp_tools/types"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; const buildTool = (schema: InputSchema | string): MCPTool => ({ name: "demo-tool", @@ -229,8 +230,7 @@ describe("ToolTestPanel argument payload", () => { const onSubmit = await submitPanel( { type: "object", properties: { active: { type: "boolean", default: false } } }, async (user) => { - await user.click(screen.getByLabelText("active")); - await user.click(await screen.findByText("True")); + await chooseSelectOption(user, screen.getByLabelText("active"), "True"); }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 8dca87e24ef..8abb8855e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -45,7 +45,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -53,11 +53,11 @@ describe("MCPServers", () => { // Wait for the component to load and check if title renders await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the title is rendered - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); it("should render mocked MCP servers data in the table", async () => { @@ -96,7 +96,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); const queryClient = createQueryClient(); - const { getByText, getAllByText } = render( + render( , @@ -104,19 +104,19 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Wait for the mocked data to render in the table await waitFor(() => { - expect(getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); }); // Verify the mocked server data is rendered in the table - expect(getByText("Test Server 1")).toBeInTheDocument(); - expect(getByText("Test Server 2")).toBeInTheDocument(); - expect(getAllByText("test-server-1").length).toBeGreaterThan(0); - expect(getAllByText("test-server-2").length).toBeGreaterThan(0); + expect(screen.getByText("Test Server 1")).toBeInTheDocument(); + expect(screen.getByText("Test Server 2")).toBeInTheDocument(); + expect(screen.getAllByText("test-server-1").length).toBeGreaterThan(0); + expect(screen.getAllByText("test-server-2").length).toBeGreaterThan(0); // Verify the API was called // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock @@ -168,7 +168,7 @@ describe("MCPServers", () => { vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -176,7 +176,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify the health check API was called (without a server ID filter — the hook always @@ -211,7 +211,7 @@ describe("MCPServers", () => { ); const queryClient = createQueryClient(); - const { getByText } = render( + render( , @@ -219,7 +219,7 @@ describe("MCPServers", () => { // Wait for the component to load await waitFor(() => { - expect(getByText("MCP Servers")).toBeInTheDocument(); + expect(screen.getByText("MCP Servers")).toBeInTheDocument(); }); // Verify that health check was initiated diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx index 8c683f230e0..f460b77c2c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx @@ -140,6 +140,7 @@ const renderPanel = (canModify = true) => accessToken="token" userRole="Admin" userID="u-admin" + isViewOnly={false} teams={null} createScope={canModify ? "unscoped-ok" : "forbidden"} />, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx index db5120cebce..5b53217f9c1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx @@ -21,12 +21,20 @@ interface AutoRoutersPanelProps { accessToken: string; userRole: string; userID: string | null; + isViewOnly: boolean; teams: Team[] | null; /** Owned by the page, which knows how this caller must scope what they create. */ createScope: ModelWriteScope; } -export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createScope }: AutoRoutersPanelProps) { +export function AutoRoutersPanel({ + accessToken, + userRole, + userID, + isViewOnly, + teams, + createScope, +}: AutoRoutersPanelProps) { const canCreate = createScope !== "forbidden"; const { data: deployments, isLoading } = useAutoRouters(); const invalidateAutoRouters = useInvalidateAutoRouters(); @@ -39,8 +47,8 @@ export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createS const [isDeleting, setIsDeleting] = useState(false); const routers = useMemo( - () => toAutoRouterRows(deployments ?? [], { userRole, userID }, teams), - [deployments, userRole, userID, teams], + () => toAutoRouterRows(deployments ?? [], { userRole, userID, isViewOnly }, teams), + [deployments, userRole, userID, isViewOnly, teams], ); const handleCreated = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts index 9944653b638..23585f6c110 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts @@ -5,8 +5,9 @@ import { toAutoRouterRow, toAutoRouterRows } from "./autoRouterRows"; // Existing cases assert resource classification, so they run as a proxy admin: the actor // gate is then a pass-through and canEdit/canDelete still reflect the row itself. -const ADMIN = { userRole: "Admin", userID: "u-admin" }; -const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" }; +const ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false }; +const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false }; +const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true }; const complexityDeployment = { model_name: "tri-tier-router", @@ -224,7 +225,7 @@ describe("autoRouterRows actor gating", () => { { team_id: "team-1", members_with_roles: [{ user_id: "u-team-admin", user_email: "t@t", role: "admin" }] }, ] as never; - const rowIn = (actor: { userRole: string; userID: string }, teamId: string | null) => + const rowIn = (actor: { userRole: string; userID: string; isViewOnly: boolean }, teamId: string | null) => toAutoRouterRow( { ...complexityDeployment, model_info: { id: "cid-1", db_model: true, team_id: teamId } }, 0, @@ -259,4 +260,12 @@ describe("autoRouterRows actor gating", () => { expect(row.canEdit).toBe(true); expect(row.canDelete).toBe(true); }); + + // A proxy_admin_viewer session reads "Admin" through the masquerade, but PATCH and + // DELETE both 403 it, so its rows must not offer the affordances. + it("hides write affordances from a view-only admin session", () => { + const row = rowIn(VIEW_ONLY_ADMIN, null); + expect(row.canEdit).toBe(false); + expect(row.canDelete).toBe(false); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx index 8b34d61ebad..7cd418bc176 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.test.tsx @@ -1,5 +1,5 @@ /* @vitest-environment jsdom */ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PriceDataManagementTab from "./PriceDataManagementTab"; @@ -11,7 +11,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ describe("PriceDataManagementTab", () => { it("renders its content standalone, without a tab-panel ancestor", () => { - const { getByText } = render(); - expect(getByText("Price Data Management")).toBeInTheDocument(); + render(); + expect(screen.getByText("Price Data Management")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 521f89a39f2..84a05113177 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsPage from "./page"; @@ -38,8 +38,17 @@ vi.mock("./useModelDashboardData", () => ({ useModelDashboardData: () => ({ availableModelAccessGroups: [], allModelsOnProxy: [], availableModelGroups: [] }), })); -const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false }; -const NON_ADMIN = { accessToken: "at", token: "t", userRole: "Internal User", userId: "u1", premiumUser: false }; +const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false, isViewOnly: false }; +const NON_ADMIN = { + accessToken: "at", + token: "t", + userRole: "Internal User", + userId: "u1", + premiumUser: false, + isViewOnly: false, +}; +// A proxy_admin_viewer session: effectiveSessionRole masquerades the role as "Admin". +const VIEW_ONLY_ADMIN = { ...ADMIN, isViewOnly: true }; const renderPage = () => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); @@ -64,48 +73,64 @@ describe("ModelsAndEndpointsPage", () => { }); it("renders the admin tab bar and the All Models panel by default", () => { - const { getByRole, getByTestId } = renderPage(); - expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); - expect(getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); - expect(getByTestId("panel-all-models")).toBeInTheDocument(); + renderPage(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "LLM Credentials" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); + expect(screen.getByTestId("panel-all-models")).toBeInTheDocument(); }); it("switches tabs in-memory, mounting only the active panel", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId, queryByTestId } = renderPage(); - await user.click(getByRole("tab", { name: "Health Status" })); - expect(getByTestId("panel-health")).toBeInTheDocument(); - expect(queryByTestId("panel-all-models")).not.toBeInTheDocument(); + renderPage(); + await user.click(screen.getByRole("tab", { name: "Health Status" })); + expect(screen.getByTestId("panel-health")).toBeInTheDocument(); + expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); }); it("renders the model detail overlay from the ?model drill-in and hides the tabs", () => { detailState.modelId = "abc-123"; - const { getByTestId, queryByRole } = renderPage(); - expect(getByTestId("model-info")).toHaveTextContent("model:abc-123"); - expect(queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.getByTestId("model-info")).toHaveTextContent("model:abc-123"); + expect(screen.queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); }); it("renders the team detail overlay from the ?team drill-in", () => { detailState.teamId = "team-9"; - const { getByTestId } = renderPage(); - expect(getByTestId("team-info")).toHaveTextContent("team:team-9"); + renderPage(); + expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); }); it("hides admin-only tabs for a non-admin user", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); - expect(queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); - expect(queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); + renderPage(); + expect(screen.queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); + }); + + // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one. + it("hides the Add Model tab for a view-only admin session", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + }); + + // Read parity: the Auto-Routers list stays reachable for a view-only admin; only the + // create affordance inside it is withheld, which AutoRoutersTabPanel decides. + it("keeps the Auto-Routers tab for a view-only admin session", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); }); // Auto-routers are excluded from the All Models table, so this tab is their home: the only // place in the product to list, create, edit or delete one. describe("Auto-Routers tab", () => { it("sits third, after All Models and Add Model", () => { - const { getAllByRole } = renderPage(); + renderPage(); - const tabs = getAllByRole("tab").map((tab) => tab.textContent); + const tabs = screen.getAllByRole("tab").map((tab) => tab.textContent); expect(tabs[0]).toContain("All Models"); expect(tabs[1]).toBe("Add Model"); expect(tabs[2]).toContain("Auto-Routers"); @@ -115,17 +140,17 @@ describe("ModelsAndEndpointsPage", () => { it("renders its panel when selected", async () => { const user = userEvent.setup(); - const { getByRole, getByTestId } = renderPage(); + renderPage(); - await user.click(getByRole("tab", { name: /Auto-Routers/ })); - expect(getByTestId("panel-auto-routers")).toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: /Auto-Routers/ })); + expect(screen.getByTestId("panel-auto-routers")).toBeInTheDocument(); }); it("is hidden from non-admins, who cannot write models", () => { mockUseAuthorized.mockReturnValue(NON_ADMIN); - const { queryByRole } = renderPage(); + renderPage(); - expect(queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: /Auto-Routers/ })).not.toBeInTheDocument(); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 9ae7dc12f81..34c9d87004e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -80,7 +80,7 @@ const renderPanel = (key: string) => { }; export default function ModelsAndEndpointsPage() { - const { accessToken, userRole, userId: userID, premiumUser } = useAuthorized(); + const { accessToken, userRole, userId: userID, premiumUser, isViewOnly } = useAuthorized(); const { data: teams } = useTeams(); const { data: uiSettings } = useUISettings(); const queryClient = useQueryClient(); @@ -92,7 +92,7 @@ export default function ModelsAndEndpointsPage() { const isInternalUser = userRole && internalUserRoles.includes(userRole); const canCreate = canCreateModels( - { userRole, userID }, + { userRole, userID, isViewOnly }, { teams: teams ?? null, disabledForInternalUsers: @@ -182,6 +182,7 @@ export default function ModelsAndEndpointsPage() { accessToken={accessToken} userID={userID} userRole={userRole} + isViewOnly={isViewOnly} onModelUpdate={invalidateModels} modelAccessGroups={availableModelAccessGroups} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx index 19e1e3aa8bd..efba26734ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.integration.test.tsx @@ -87,7 +87,6 @@ const alwaysMounted = { api_key: undefined, api_base: undefined, custom_llm_provider: "openai", - litellm_credential_name: null, model: "gpt-4o", }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx new file mode 100644 index 00000000000..12f0b95bf13 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx @@ -0,0 +1,39 @@ +/* @vitest-environment jsdom */ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import AutoRoutersTabPanel from "./AutoRoutersTabPanel"; + +const panelProps = vi.fn(); +vi.mock("../components/AutoRouters/AutoRoutersPanel", () => ({ + AutoRoutersPanel: (props: Record) => { + panelProps(props); + return
; + }, +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized() })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => ({ data: [] }) })); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: () => ({ data: { values: {} } }), +})); + +const SESSION = { accessToken: "at", userRole: "Admin", userId: "u1", isViewOnly: false }; + +const lastProps = () => panelProps.mock.calls.at(-1)?.[0] as { createScope: string }; + +describe("AutoRoutersTabPanel", () => { + it("grants an unscoped create to a real proxy admin", () => { + mockUseAuthorized.mockReturnValue(SESSION); + render(); + expect(lastProps().createScope).toBe("unscoped-ok"); + }); + + // The masqueraded "Admin" a proxy_admin_viewer session carries: POST /model/new 403s it, + // so the panel must not be told it may create. + it("withholds the create affordance from a view-only admin session", () => { + mockUseAuthorized.mockReturnValue({ ...SESSION, isViewOnly: true }); + render(); + expect(lastProps().createScope).toBe("forbidden"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx index 5f7d56e8e33..69b442da09b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx @@ -15,13 +15,13 @@ import { AutoRoutersPanel } from "../components/AutoRouters/AutoRoutersPanel"; * Viewer roles reach the list without write affordances. */ export default function AutoRoutersTabPanel() { - const { accessToken, userRole, userId: userID } = useAuthorized(); + const { accessToken, userRole, userId: userID, isViewOnly } = useAuthorized(); const { data: teams } = useTeams(); const { data: uiSettings } = useUISettings(); const isInternalUser = userRole != null && internalUserRoles.includes(userRole); const scope = modelCreationScope( - { userRole, userID }, + { userRole, userID, isViewOnly }, { teams: teams ?? null, disabledForInternalUsers: isInternalUser && uiSettings?.values?.disable_model_add_for_internal_users === true, @@ -33,6 +33,7 @@ export default function AutoRoutersTabPanel() { accessToken={accessToken} userRole={userRole ?? ""} userID={userID ?? null} + isViewOnly={isViewOnly} teams={teams ?? null} createScope={scope} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index a83c11d1444..ed02e7e16cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -93,13 +93,13 @@ describe("ChatMessageBubble", () => { ])("should paint the $role surface from theme tokens, not fixed colours", ({ role, bubble, avatar }) => { render(); - const header = screen.getByText(role).closest("div") as HTMLElement; - const surface = header.parentElement as HTMLElement; + const surface = screen.getByTestId("message-surface"); + const avatarEl = screen.getByTestId("message-avatar"); expect(surface).toHaveClass(...bubble); expect(surface).not.toHaveAttribute("style"); - expect(header.firstElementChild).toHaveClass(avatar); - expect(header.firstElementChild).not.toHaveAttribute("style"); + expect(avatarEl).toHaveClass(avatar); + expect(avatarEl).not.toHaveAttribute("style"); }); it("should show model badge for assistant messages when model is provided", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index 8c54d9e89fa..bd6a4bc49a4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -46,6 +46,7 @@ function ChatMessageBubble({ return (
{ describe("CompareUI", () => { it("should render", () => { - const { getByTestId } = render(); - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); - expect(getByTestId("message-input")).toBeInTheDocument(); + render(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("message-input")).toBeInTheDocument(); }); it("adds a comparison when Add Comparison button is clicked", async () => { const user = userEvent.setup(); - const { container, getByTestId } = render( - , - ); + const { container } = render(); // Verify initial state: 2 comparison panels - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); let comparisonPanels = container.querySelectorAll('[data-testid^="comparison-panel-"]'); expect(comparisonPanels).toHaveLength(2); @@ -117,15 +115,13 @@ describe("CompareUI", () => { }); // Verify the original 2 panels are still there - expect(getByTestId("comparison-panel-1")).toBeInTheDocument(); - expect(getByTestId("comparison-panel-2")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-1")).toBeInTheDocument(); + expect(screen.getByTestId("comparison-panel-2")).toBeInTheDocument(); }); it("should handle image upload and send message with attachment", async () => { const user = userEvent.setup(); - const { getByTestId, queryByTestId } = render( - , - ); + render(); const file = new File(["test content"], "test-image.png", { type: "image/png" }); @@ -138,13 +134,13 @@ describe("CompareUI", () => { } await waitFor(() => { - expect(getByTestId("has-attachment")).toBeInTheDocument(); + expect(screen.getByTestId("has-attachment")).toBeInTheDocument(); }); - const textarea = getByTestId("message-textarea"); + const textarea = screen.getByTestId("message-textarea"); fireEvent.change(textarea, { target: { value: "Describe this image" } }); - const sendButton = getByTestId("send-button"); + const sendButton = screen.getByTestId("send-button"); expect(sendButton).toBeEnabled(); await user.click(sendButton); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx index 72a1d41f9fe..69b50d4088f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import type { MessageType } from "@/components/chat_ui/types"; import { MessageDisplay } from "./MessageDisplay"; @@ -39,9 +39,9 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByText } = render(); - expect(getByText("Hello")).toBeInTheDocument(); - expect(getByText("Hi there!")).toBeInTheDocument(); + render(); + expect(screen.getByText("Hello")).toBeInTheDocument(); + expect(screen.getByText("Hi there!")).toBeInTheDocument(); }); it("displays user and assistant messages with proper grouping and shows loading state", () => { @@ -64,13 +64,13 @@ describe("MessageDisplay", () => { }, }, ]; - const { getByText, getByTestId } = render(); - expect(getByText("You")).toBeInTheDocument(); - expect(getByText("What is 2+2?")).toBeInTheDocument(); - expect(getByText("gpt-4")).toBeInTheDocument(); - expect(getByText("calculator")).toBeInTheDocument(); - expect(getByText("2+2 equals 4")).toBeInTheDocument(); - expect(getByTestId("response-metrics")).toBeInTheDocument(); + render(); + expect(screen.getByText("You")).toBeInTheDocument(); + expect(screen.getByText("What is 2+2?")).toBeInTheDocument(); + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("calculator")).toBeInTheDocument(); + expect(screen.getByText("2+2 equals 4")).toBeInTheDocument(); + expect(screen.getByTestId("response-metrics")).toBeInTheDocument(); }); it("should display image attachment in user message", () => { @@ -86,10 +86,10 @@ describe("MessageDisplay", () => { model: "gpt-4", }, ]; - const { getByTestId, getByText } = render(); - expect(getByText("What is in this image? [Image attached]")).toBeInTheDocument(); - expect(getByTestId("chat-image-renderer")).toBeInTheDocument(); - const image = getByTestId("chat-image-renderer").querySelector("img"); + render(); + expect(screen.getByText("What is in this image? [Image attached]")).toBeInTheDocument(); + expect(screen.getByTestId("chat-image-renderer")).toBeInTheDocument(); + const image = screen.getByTestId("chat-image-renderer").querySelector("img"); expect(image).toHaveAttribute("src", "blob:test-image-url"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx index c8779be06c4..f23145a917a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/policy_test_panel.test.tsx @@ -5,6 +5,7 @@ import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event" import { renderWithProviders } from "@/../tests/test-utils"; import * as networking from "@/components/networking"; import PolicyTestPanel from "./policy_test_panel"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; vi.mock("@/components/networking"); @@ -24,8 +25,7 @@ const setup = () => { }; const pickOption = async (user: ReturnType, label: string, option: string) => { - await user.click(screen.getByLabelText(label)); - await user.click(await screen.findByTitle(option)); + await chooseSelectOption(user, screen.getByLabelText(label), option); }; const simulate = async (user: ReturnType) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx index d6a4aaea2e1..3b0ee3a3ce2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { deletePromptCall, getPromptsList } from "@/components/networking"; import PromptsPanel from "./index"; +import { chooseSelectOption } from "../../../../../tests/test-utils"; vi.mock("@/components/networking", () => ({ getPromptsList: vi.fn(), @@ -118,8 +119,7 @@ describe("PromptsPanel toolbar", () => { expect(screen.getByText("All Environments")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Production")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Production"); await waitFor(() => expect(mockGetPromptsList).toHaveBeenLastCalledWith("sk-test", "production")); }); @@ -131,12 +131,10 @@ describe("PromptsPanel toolbar", () => { renderPanel("Admin"); await screen.findByText("table-loaded"); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Production")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Production"); await waitFor(() => expect(screen.getByRole("combobox")).toHaveTextContent("Production")); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("All Environments")); + await chooseSelectOption(user, screen.getByRole("combobox"), "All Environments"); await waitFor(() => expect(screen.getByRole("combobox")).toHaveTextContent("All Environments")); await waitFor(() => expect(mockGetPromptsList).toHaveBeenLastCalledWith("sk-test", undefined)); diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/ui/litellm-dashboard/src/autorouter_presets.json index 4cbb548a855..b977bd484ac 100644 --- a/ui/litellm-dashboard/src/autorouter_presets.json +++ b/ui/litellm-dashboard/src/autorouter_presets.json @@ -14,7 +14,9 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -30,7 +32,9 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -56,7 +60,9 @@ }, "classifier_context_window_size": 0, "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } }, @@ -72,7 +78,9 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], + "classification_mode": "every_request", "session_affinity": false, + "modality_routing": false, "deployment_affinity": true } } diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index 397e2ab10fa..906565750b1 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -203,12 +203,12 @@ describe("ModelHubTable", () => { }); describe("hub tabs", () => { - const renderHub = async () => { + const renderHub = async (agents: object[] = []) => { vi.mocked(networking.modelHubCall).mockResolvedValue({ data: [{ model_group: "claude-opus-4-8", providers: ["anthropic"], mode: "chat" }], }); vi.mocked(networking.getConfigFieldSetting).mockResolvedValue({ field_value: false }); - vi.mocked(networking.getAgentsList).mockResolvedValue({ agents: [] }); + vi.mocked(networking.getAgentsList).mockResolvedValue({ agents }); vi.mocked(networking.fetchMCPServers).mockResolvedValue([]); vi.mocked(networking.getUiSettings).mockResolvedValue({ values: {} }); mockUseUISettings.mockReturnValue({ data: { values: {} }, isLoading: false }); @@ -230,6 +230,37 @@ describe("ModelHubTable", () => { expect(await screen.findByPlaceholderText("Search model names...")).toHaveValue("opus"); }); + it("filters the Agent Hub table by name or description and shows the no-match state", async () => { + const { user } = await renderHub([ + { + agent_id: "a1", + agent_card_params: { name: "Billing Router", description: "routes billing questions" }, + litellm_params: { is_public: false }, + }, + { + agent_id: "a2", + agent_card_params: { name: "Support Bot", description: "handles support tickets" }, + litellm_params: { is_public: false }, + }, + ]); + const agentCount = (expected: string) => + screen.getByText((_, el) => el?.tagName === "P" && el.textContent === expected); + + await user.click(screen.getByRole("tab", { name: "Agent Hub" })); + expect(await screen.findByText("Billing Router")).toBeInTheDocument(); + + const search = screen.getByPlaceholderText("Search agent names or descriptions..."); + await user.type(search, "support tickets"); + expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); + expect(screen.getByText("Support Bot")).toBeInTheDocument(); + expect(agentCount("Showing 1 of 2 agents")).toBeInTheDocument(); + + await user.clear(search); + await user.type(search, "zzzz"); + expect(screen.getByText("No matching agents")).toBeInTheDocument(); + expect(agentCount("Showing 0 of 2 agents")).toBeInTheDocument(); + }); + it("renders the hub strip as underlined tabs rather than a segmented pill", async () => { await renderHub(); diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 299104b271b..475e3dcd70b 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -23,13 +23,15 @@ import { import PublicModelHub from "@/components/public_model_hub"; import { copyToClipboard } from "@/utils/dataUtils"; import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; +import { filterBySearchTerm } from "@/utils/searchUtils"; import { SortingState } from "@tanstack/react-table"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { Copy, Inbox } from "lucide-react"; +import { Copy, Inbox, Search as SearchIcon, X } from "lucide-react"; import { useRouter } from "next/navigation"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; @@ -80,6 +82,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, const [agentLoading, setAgentLoading] = useState(true); const [selectedAgent, setSelectedAgent] = useState(null); const [isAgentModalVisible, setIsAgentModalVisible] = useState(false); + const [agentSearchTerm, setAgentSearchTerm] = useState(""); // MCP Hub state const [mcpHubData, setMcpHubData] = useState(null); const [mcpLoading, setMcpLoading] = useState(true); @@ -385,6 +388,10 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, const modelColumns = useMemo(() => getModelHubTableColumns({ onModelClick: showModal }), [showModal]); const agentColumns = useMemo(() => getAgentHubTableColumns({ onAgentClick: showAgentModal }), [showAgentModal]); + const filteredAgentData = useMemo( + () => filterBySearchTerm(agentHubData ?? [], agentSearchTerm, (agent) => [agent.name, agent.description]), + [agentHubData, agentSearchTerm], + ); const mcpColumns = useMemo(() => getMCPHubTableColumns({ onServerClick: showMcpModal }), [showMcpModal]); // If this is a public page, use the dedicated PublicModelHub component @@ -505,9 +512,34 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage,
)} +
+

Search Agents:

+ + + + + setAgentSearchTerm(e.target.value)} + /> + {agentSearchTerm && ( + + setAgentSearchTerm("")} + > + + + + )} + +
+ {/* Agent Table */} agent.agent_id || agent.name || String(index)} sortingMode="client" @@ -516,7 +548,14 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, isLoading={agentLoading} loadingMessage="Loading agents…" noDataMessage={ - + } size="compact" /> @@ -524,7 +563,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage,

- Showing {agentHubData?.length || 0} agent{agentHubData?.length !== 1 ? "s" : ""} + Showing {filteredAgentData.length} of {agentHubData?.length || 0} agents

diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 3777c46973f..68bfd94bd63 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -5,6 +5,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { CreateUserButton } from "./CreateUserButton"; import * as networking from "./networking"; import { toast } from "@/lib/toast"; +import { expectControlBesideLabel } from "../../tests/fieldOrientation"; vi.mock("./networking", () => ({ userCreateCall: vi.fn(), @@ -294,6 +295,20 @@ describe("CreateUserButton", () => { }); }); + it("lays the send invitation email checkbox out beside its label", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + + expectControlBesideLabel(within(dialog).getByRole("checkbox")); + }); + describe("organizations", () => { it("should send organizations list in POST body when organizations are selected", async () => { const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.tsx index e052a9e0818..0f7c356b8cc 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.tsx @@ -270,7 +270,7 @@ export const CreateUserButton: React.FC = ({ ); const sendInviteEmailField = ( - + {({ id, value, onChange, onBlur }) => ( )} diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx index 8edf2174eee..cb04323e4c9 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx @@ -10,6 +10,7 @@ */ import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import EntityUsageExportModal from "./EntityUsageExportModal"; @@ -73,13 +74,13 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByRole } = renderWithProviders(); + renderWithProviders(); // Default primary action reflects CSV export - expect(getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /Export CSV/i })).toBeInTheDocument(); // Click export - await user.click(getByRole("button", { name: /Export CSV/i })); + await user.click(screen.getByRole("button", { name: /Export CSV/i })); // Verifies export function was invoked with correct parameters expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily", "Tag", "tag", {}); @@ -97,14 +98,14 @@ describe("EntityUsageExportModal", () => { const user = userEvent.setup(); const { handleExportCSV } = await import("./utils"); - const { getByText, getByRole } = renderWithProviders(); + renderWithProviders(); // Choose the alternate export type - click the label to trigger radio - const dailyModelLabel = getByText(/Day-by-day by tag and model/i); + const dailyModelLabel = screen.getByText(/Day-by-day by tag and model/i); await user.click(dailyModelLabel); // Export with default CSV format - const exportBtn = getByRole("button", { name: /Export CSV/i }); + const exportBtn = screen.getByRole("button", { name: /Export CSV/i }); await user.click(exportBtn); // Ensure the selected scope flowed through diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx index 54566e1d75a..eddc5716305 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.test.tsx @@ -9,6 +9,7 @@ import BaseSSOSettingsForm, { submitMountedSSOValues, useSSOSettingsForm, } from "./BaseSSOSettingsForm"; +import { expectControlBesideLabel } from "../../../../../../tests/fieldOrientation"; const user = () => userEvent.setup({ pointerEventsCheck: 0 }); @@ -233,6 +234,38 @@ describe("BaseSSOSettingsForm", () => { expect(screen.queryByText("Use Team Mappings")).not.toBeInTheDocument(); }); + + it("lays a provider checkbox field out beside its label", async () => { + const TestWrapper = () => { + const form = useSSOSettingsForm("sso-settings"); + + return ; + }; + + renderWithProviders(); + + await openProviderDropdown(); + await user().click(await screen.findByText(/saml sso/i)); + + expectControlBesideLabel( + await screen.findByRole("checkbox", { name: "Allow IdP-initiated (unsolicited) responses" }), + ); + }); + + it.each(["Use Role Mappings", "Use Team Mappings"])("lays the %s toggle out beside its label", async (label) => { + const TestWrapper = () => { + const form = useSSOSettingsForm("sso-settings"); + + return ; + }; + + renderWithProviders(); + + await openProviderDropdown(); + await user().click(await screen.findByText(/okta/i)); + + expectControlBesideLabel(await screen.findByRole("checkbox", { name: label })); + }); }); describe("renderProviderFields", () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx index 5216da382d0..cb97304f77a 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -303,7 +303,7 @@ const SSOProviderField = ({ field }: { field: SSOProviderConfig["fields"][number if (field.type === "checkbox") { return ( - + {({ value, onChange, onBlur, id, ...rest }) => ( (); return ( - + {({ value, onChange, onBlur, id, ...rest }) => ( { await user.click(await screen.findByText("Edit team")); expect(onEditTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); - await user.click(screen.getByTestId("team-actions-team-1")); - await user.click(await screen.findByText("Delete team")); + await chooseSelectOption(user, screen.getByTestId("team-actions-team-1"), "Delete team", "menuitem"); expect(onDeleteTeam).toHaveBeenCalledWith(expect.objectContaining({ team_id: "team-1" })); }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index d09ccbf8826..b06448912d0 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -2,7 +2,7 @@ import { screen, waitFor, within, fireEvent } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; @@ -327,8 +327,7 @@ it("sorts by the backend max_budget field when 'Budget descending' is chosen fro const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith( @@ -343,8 +342,7 @@ it("emphasizes the active field in the Spend / Budget header so the sorted colum const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); await waitFor(() => { expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" })).toHaveClass("font-semibold"); @@ -356,8 +354,7 @@ it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / B const user = userEvent.setup(); renderWithProviders(); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Spend ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Spend ascending", "menuitem"); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "spend", sortOrder: "asc" })); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index ad0f749b189..92951b68fd5 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -82,7 +82,7 @@ const AddModelForm: React.FC = ({ // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test const [connectionTestId, setConnectionTestId] = useState(""); - const { accessToken, userRole, premiumUser, userId } = useAuthorized(); + const { accessToken, userRole, premiumUser, userId, isViewOnly } = useAuthorized(); const { data: providerMetadata, isLoading: isProviderMetadataLoading, @@ -157,7 +157,10 @@ const AddModelForm: React.FC = ({ const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); // Same owner the Auto-Routers tab uses, so the two creation forms cannot disagree about // who has to name a team. This form is only reachable when creation is allowed at all. - const createScope = modelCreationScope({ userRole, userID: userId }, { teams, disabledForInternalUsers: false }); + const createScope = modelCreationScope( + { userRole, userID: userId, isViewOnly }, + { teams, disabledForInternalUsers: false }, + ); const requiresTeamScope = createScope === "team-required"; return ( diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx new file mode 100644 index 00000000000..4d9e2122739 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -0,0 +1,26 @@ +import React from "react"; + +import { Switch } from "@/components/ui/switch"; + +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig"; + +export const AffinityControls: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => ( + <> +
+ onChange({ ...value, deployment_affinity: deploymentAffinity })} + aria-label="Pin a session to one deployment per model group" + /> + Pin a session to one deployment per model group +
+ + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to + load-balance every turn. + + +); diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 84f61c95eca..96c93306611 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -15,9 +15,12 @@ import { RestrictedSection, restrictedBy } from "./TierRestrictions"; import HeuristicScoringConfig from "./HeuristicScoringConfig"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { + ClassificationFrequency, ClassifierFallback, ClassifierType, ComplexityRouterConfigValue, + classificationFrequency, + withClassificationFrequency, DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, MIN_QUOTED_CONTEXT_TURN_CHARS, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, @@ -211,6 +214,7 @@ const ClassificationMethodConfig: React.FC = ({ const [draft, setDraft] = React.useState<{ id: string; raw: string } | null>(null); const hasDefaultModel = Boolean(defaultModel); const classifierType = effectiveClassifierType(value); + const sessionFrequencyRestriction = restrictedBy(value, "sessionAffinity"); const classifierModelMissing = showValidationErrors && usesLlmClassifier(classifierType) && !value.classifier_llm_config?.model; const usesCustomPrompt = Boolean(value.classifier_llm_config?.system_prompt?.trim()); @@ -305,6 +309,10 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, classifier_fallback: fallback }); }; + const handleClassificationFrequencyChange = (frequency: ClassificationFrequency) => { + onChange(withClassificationFrequency(value, frequency)); + }; + const handleClassifierContextWindowSizeChange = (windowSize: number) => { onChange({ ...value, @@ -367,6 +375,49 @@ const ClassificationMethodConfig: React.FC = ({
)} +
+ How often to classify + + handleClassificationFrequencyChange(frequency as ClassificationFrequency) + } + > +
+ + + +
+ +

+ Holding the tier keeps an agent on one model for a whole tool loop and cuts scoring cost. A turn the router + cannot match to a held decision, such as one with no session id or an expired one, is scored again +

+
+ {usesLlmClassifier(classifierType) && (
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 218c99e32c5..751f8870561 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -598,6 +598,82 @@ describe("ComplexityRouterConfig classifier fallback", () => { }); }); +describe("ComplexityRouterConfig classification frequency", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + + it("defaults to every request, matching both backend field defaults", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every request/ })).toBeChecked(); + expect(screen.getByRole("radio", { name: /Every new user message/ })).not.toBeChecked(); + expect(screen.getByRole("radio", { name: /Once per session/ })).not.toBeChecked(); + }); + + it("writes both wire fields when the frequency moves to every new user message", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByRole("radio", { name: /Every new user message/ })); + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classification_mode: "user_turn", + session_affinity: false, + }); + }); + + it("writes session affinity, not a classification mode, when the frequency moves to once per session", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByRole("radio", { name: /Once per session/ })); + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classification_mode: "every_request", + session_affinity: true, + }); + }); + + it("shows a hand-authored config that sets both fields as once per session, matching the backend", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Once per session/ })).toBeChecked(); + expect(screen.getByRole("radio", { name: /Every new user message/ })).not.toBeChecked(); + }); + + it("records a switch back to every request", () => { + const onChange = vi.fn(); + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every new user message/ })).toBeChecked(); + fireEvent.click(screen.getByRole("radio", { name: /Every request/ })); + expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ classification_mode: "every_request" })); + }); + + it("offers the frequency on a heuristic router, where holding the tier still pins the model", () => { + // The backend pin is gated on the two fields alone, so a heuristic router that switches models + // mid tool loop is fixed by this control too. + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.getByRole("radio", { name: /Every new user message/ })).toBeInTheDocument(); + }); +}); + describe("ComplexityRouterConfig classifier rubric", () => { const llmValue: ComplexityRouterConfigValue = { ...defaultValue, @@ -760,13 +836,34 @@ describe("ComplexityRouterConfig tier labels", () => { }); }); +describe("ComplexityRouterConfig modality panel", () => { + it("defaults the image-routing switch off and writes modality_routing through onChange", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + const toggle = screen.getByRole("switch", { name: "Route image requests to vision-capable models" }); + expect(toggle).not.toBeChecked(); + fireEvent.click(toggle); + + expect(onChange).toHaveBeenCalledWith({ ...defaultValue, modality_routing: true }); + }); + + it("renders a stored modality_routing=true as on", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + expect(screen.getByRole("switch", { name: "Route image requests to vision-capable models" })).toBeChecked(); + }); +}); + describe("ComplexityRouterConfig affinity panel", () => { - it("holds both affinity switches with their backend defaults", () => { + it("holds the deployment switch at its backend default, session pinning having moved to the frequency choice", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Affinity")); expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).toBeChecked(); - expect(screen.getByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + expect(screen.queryByRole("switch", { name: "Pin a session to its first model" })).not.toBeInTheDocument(); }); it("writes deployment_affinity through onChange without touching other keys", () => { @@ -1291,10 +1388,18 @@ describe("ComplexityRouterConfig tier editing", () => { expect(screen.getByLabelText("Fallback tier")).toBeInTheDocument(); }); - it("disables session pinning and says why, rather than letting a stripped value look saved", () => { - renderWithProviders(); - fireEvent.click(screen.getByText("Advanced: Affinity")); - expect(screen.getByLabelText("Pin a session to its first model")).toHaveAttribute("data-disabled"); + it("disables the once-per-session frequency and says why, rather than letting a stripped value look saved", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + const sessionOption = screen.getByRole("radio", { name: /Once per session/ }); + expect(sessionOption).toHaveAttribute("aria-disabled", "true"); + expect(sessionOption).not.toBeChecked(); expect( screen.getByText("Session pinning escalates along the built-in tier ladder", { exact: false }), ).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 111a7c9f10a..6062705aa82 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -4,6 +4,9 @@ import { SearchSelect } from "@/components/shared/SearchSelect"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; import { Switch } from "@/components/ui/switch"; + +import { AffinityControls } from "./AffinityControls"; +import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; @@ -56,6 +59,16 @@ export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120; export const DEFAULT_SESSION_AFFINITY = false; export const DEFAULT_DEPLOYMENT_AFFINITY = true; +export type ClassificationMode = "every_request" | "user_turn"; + +export const DEFAULT_CLASSIFICATION_MODE: ClassificationMode = "every_request"; + +/** + * One operator-facing choice over the two wire fields that share the router's tier-pin machinery: + * session affinity pins every turn, user_turn pins every turn except a new human ask. + */ +export type ClassificationFrequency = ClassificationMode | "session"; + export type ComplexityTiers = { SIMPLE: string[]; MEDIUM: string[]; @@ -384,7 +397,9 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + classification_mode?: ClassificationMode; session_affinity?: boolean; + modality_routing?: boolean; deployment_affinity?: boolean; /** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */ plan_mode_min_tier?: string; @@ -420,6 +435,21 @@ export interface ComplexityRouterConfigValue { tier_model_params?: TierModelParamsByTier; } +/** Session affinity wins where a hand-authored config sets both, matching the backend's own `or`. */ +export const classificationFrequency = (value: ComplexityRouterConfigValue): ClassificationFrequency => { + if (!value.custom_tier_set && (value.session_affinity ?? DEFAULT_SESSION_AFFINITY)) return "session"; + return value.classification_mode === "user_turn" ? "user_turn" : "every_request"; +}; + +export const withClassificationFrequency = ( + value: ComplexityRouterConfigValue, + frequency: ClassificationFrequency, +): ComplexityRouterConfigValue => ({ + ...value, + classification_mode: frequency === "user_turn" ? "user_turn" : "every_request", + session_affinity: frequency === "session", +}); + interface ComplexityRouterConfigProps { modelInfo: ModelGroup[]; value: ComplexityRouterConfigValue; @@ -485,39 +515,6 @@ export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; */ export const HEURISTIC_FIRST_MAX_TIER_KEYS = TIER_KEYS.slice(0, -1); -const AffinityControls: React.FC<{ - value: ComplexityRouterConfigValue; - onChange: (value: ComplexityRouterConfigValue) => void; -}> = ({ value, onChange }) => ( - <> -
- onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" - /> - Pin a session to one deployment per model group -
- - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. - -
- onChange({ ...value, session_affinity: sessionAffinity })} - aria-label="Pin a session to its first model" - /> - Pin a session to its first model -
- - {restrictedBy(value, "sessionAffinity")?.reason ?? - "Keeps a session on its first turn's model instead of re-classifying each turn. Also pins the deployment."} - - -); - const PlanModeOverrideControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; @@ -828,6 +825,11 @@ const ComplexityRouterConfig: React.FC = ({ label: Advanced: Affinity, children: , }, + { + key: "modality", + label: Advanced: Modality Routing, + children: , + }, { key: "plan-mode", label: Advanced: Plan-Mode Override, diff --git a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx new file mode 100644 index 00000000000..dd697b35239 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx @@ -0,0 +1,26 @@ +import React from "react"; + +import { Switch } from "@/components/ui/switch"; + +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +export const ModalityRoutingControls: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => ( + <> +
+ onChange({ ...value, modality_routing: modalityRouting })} + aria-label="Route image requests to vision-capable models" + /> + Route image requests to vision-capable models +
+ + Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default + model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are replaced, + and a kept session pin still wins. + + +); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 71b454dbb06..ba380662403 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -101,18 +101,24 @@ vi.mock("./build_complexity_router_config", async (importOriginal) => { }); // A real TeamDropdown fetches teams and renders an antd Select; the wiring under test is -// whether team_id is registered, validated and forwarded, so a plain control stands in. +// whether team_id is registered, validated and forwarded, so a plain control stands in. The +// clear button mirrors the real dropdown's x, which emits null rather than a string. vi.mock("../common_components/team_dropdown", () => ({ - default: ({ value, onChange }: { value?: string; onChange?: (next: string) => void }) => ( - + default: ({ value, onChange }: { value?: string; onChange?: (next: string | null) => void }) => ( + <> + + + ), })); @@ -354,6 +360,25 @@ describe("AddAutoRouterTab", () => { expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); }); + // The shared dropdown emits null on clear while this form's schema wants a string, so the + // form maps null back to "": the user sees the pick-a-team message, not a zod type error. + it("treats a team picked and then cleared like no team at all", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders( + , + ); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); + await user.click(screen.getByTestId("team-dropdown-clear")); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + expect(await screen.findByText("Please select a team to continue")).toBeInTheDocument(); + expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); + }); + it("defaults a new router to session affinity off, matching the backend field default", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); @@ -362,8 +387,8 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + await user.click(screen.getByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Once per session/ })).not.toBeChecked(); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -468,8 +493,8 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Once per session/ })); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -479,6 +504,44 @@ describe("AddAutoRouterTab", () => { }); }); + it("carries every new user message through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "user-turn-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + classification_mode: "user_turn", + }); + }); + + it("writes every_request into the create payload when the default frequency stays selected", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "default-timing-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every request/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect( + vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config.classification_mode, + ).toBe("every_request"); + }); + it("defaults a new router to deployment affinity on, matching the backend field default", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index ed584a4882b..03b7432e4ec 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, classifierLlmConfig: complexityRouterConfig.classifier_llm_config, @@ -350,6 +351,7 @@ const AddAutoRouterTab: React.FC = ({ classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns, classifierFallback: complexityRouterConfig.classifier_fallback, sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, + modalityRouting: complexityRouterConfig.modality_routing ?? false, deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords, keywordTierRules, @@ -546,7 +548,9 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => } + {({ id, value, onChange }) => ( + onChange(next ?? "")} /> + )} )} diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 01e00d903aa..ceb9d1ca713 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -1,4 +1,4 @@ -import { act, fireEvent, render, waitFor } from "@testing-library/react"; +import { act, fireEvent, render, waitFor, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { MountedFormHost } from "../../../tests/mounted-form-host"; import AdvancedSettings from "./advanced_settings"; @@ -35,51 +35,51 @@ describe("AdvancedSettings", () => { }); it("should render tags list", async () => { - const { getByText } = renderAdvancedSettings(); - fireEvent.click(getByText("Advanced Settings")); + renderAdvancedSettings(); + fireEvent.click(screen.getByText("Advanced Settings")); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); }); it("should render the litellm params", async () => { - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("LiteLLM Params")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Params")).toBeInTheDocument(); }); }); it("hides every PTU field when PTU cost attribution is disabled", async () => { - const { getByText, queryByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("Tags")).toBeInTheDocument(); + expect(screen.getByText("Tags")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(queryByText(label)).not.toBeInTheDocument(); + expect(screen.queryByText(label)).not.toBeInTheDocument(); } - expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); }); it("shows every PTU field when PTU cost attribution is enabled", async () => { mockUsePtuCostAttributionEnabled.mockReturnValue(true); - const { getByText } = renderAdvancedSettings(); + renderAdvancedSettings(); act(() => { - fireEvent.click(getByText("Advanced Settings")); + fireEvent.click(screen.getByText("Advanced Settings")); }); await waitFor(() => { - expect(getByText("PTU Count")).toBeInTheDocument(); + expect(screen.getByText("PTU Count")).toBeInTheDocument(); }); for (const label of PTU_LABELS) { - expect(getByText(label)).toBeInTheDocument(); + expect(screen.getByText(label)).toBeInTheDocument(); } - expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index feddaaa0eac..af87cc6c8fb 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -52,8 +52,10 @@ describe("buildComplexityRouterConfig", () => { const expected = { tiers, classifier_type: "heuristic", + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, escalation_keywords: ["LITELLM ESCALATE"], }; expect(config).toEqual(expected); @@ -247,6 +249,12 @@ describe("buildComplexityRouterConfig", () => { expect(config.return_raw_model_name).toBeUndefined(); }); + it("writes modality_routing explicitly both ways, so the stored config never relies on the backend default", () => { + expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: true }).modality_routing).toBe(true); + expect(buildComplexityRouterConfig(baseParams).modality_routing).toBe(false); + expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: false }).modality_routing).toBe(false); + }); + it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => { const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true }); expect(config.session_affinity).toBe(true); @@ -790,6 +798,20 @@ describe("heuristic_first", () => { }); }); +describe("classification_mode", () => { + it("emits user_turn", () => { + const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); + expect(config.classification_mode).toBe("user_turn"); + }); + + it("writes every_request explicitly, so a saved router never depends on the backend default", () => { + expect( + buildComplexityRouterConfig({ ...baseParams, classificationMode: "every_request" }).classification_mode, + ).toBe("every_request"); + expect(buildComplexityRouterConfig(baseParams).classification_mode).toBe("every_request"); + }); +}); + describe("buildComplexityRouterConfig with an edited tier set", () => { const customTierSet = { tiers: [ @@ -902,6 +924,10 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { expect(payload.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); }); + it("keeps classification_mode, which the backend accepts beside tier_definitions", () => { + expect(build({ classificationMode: "user_turn" }).classification_mode).toBe("user_turn"); + }); + it("carries the plan-mode floor as the row's name, not the row id the form holds", () => { expect(build({ planModeMinTier: "sec" }).plan_mode_min_tier).toBe("SECURITY_REVIEW"); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 9a14e956207..af34dc92c0f 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -19,10 +19,12 @@ import { import { AdaptiveEligible, AdaptiveRouterWeights, + ClassificationMode, ClassifierFallback, ClassifierLLMConfig, ClassifierType, ComplexityTierLabels, + DEFAULT_CLASSIFICATION_MODE, ComplexityRouterConfigValue, ComplexityTiers, DimensionWeights, @@ -105,7 +107,9 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; + modalityRouting?: boolean; deploymentAffinity: boolean; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; @@ -159,8 +163,10 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; + modality_routing: boolean; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -393,7 +399,9 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + classificationMode, sessionAffinity, + modalityRouting, deploymentAffinity, customTechnicalKeywords, keywordTierRules, @@ -452,8 +460,10 @@ export const buildComplexityRouterConfig = ({ ...(cleanedTierLabels && { tier_labels: cleanedTierLabels }), classifier_type: classifierType, ...classifierWireFields(effectiveType, classifierInputs), + classification_mode: classificationMode ?? DEFAULT_CLASSIFICATION_MODE, session_affinity: sessionAffinity, deployment_affinity: deploymentAffinity, + modality_routing: modalityRouting ?? false, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 7ef09d34924..9d792480c9f 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -1,6 +1,10 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { prepareModelAddRequest } from "./handle_add_model_submit"; +vi.mock("../networking", () => ({ + modelCreateCall: vi.fn(), +})); + describe("prepareModelAddRequest", () => { it("returns deployment data for the most basic form", async () => { const formValues = { @@ -73,4 +77,28 @@ describe("prepareModelAddRequest", () => { expect(deployment.litellmParamsObj.litellm_credential_name).toBe("selected-credential"); expect(deployment.litellmParamsObj.timeout).toBe(5); }); + + it("keeps litellm_credential_name from LiteLLM Params JSON when no credential is selected", async () => { + const formValues = { + model_mappings: [ + { + public_name: "Public Model", + litellm_model: "litellm/public", + }, + ], + model_name: "custom-model-name", + litellm_extra_params: JSON.stringify({ + litellm_credential_name: "from-json", + timeout: 5, + }), + litellm_credential_name: null, + }; + + const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); + + expect(deployments).toHaveLength(1); + const [deployment] = deployments!; + expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json"); + expect(deployment.litellmParamsObj.timeout).toBe(5); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index bb2f78fa84e..41133958c0a 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -91,6 +91,9 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value === "") { continue; } + if (key === "litellm_credential_name" && value == null) { + continue; + } // Skip the custom_pricing and pricing_model fields as they're only used for UI control if (key === "custom_pricing" || key === "pricing_model" || key === "cache_control") { continue; @@ -124,13 +127,13 @@ export const prepareModelAddRequest = async (formValues: Record, ac if (value && value != undefined) { try { litellmExtraParams = JSON.parse(value); - if ("litellm_credential_name" in litellmExtraParams) { - delete litellmExtraParams.litellm_credential_name; - } } catch (error) { toast.fromError("Failed to parse LiteLLM Extra Params: " + error); throw new Error("Failed to parse litellm_extra_params: " + error); } + if ("litellm_credential_name" in litellmExtraParams && formValues.litellm_credential_name) { + delete litellmExtraParams.litellm_credential_name; + } for (const [key, value] of Object.entries(litellmExtraParams)) { litellmParamsObj[key] = value; } diff --git a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx index 64074481603..d6476089cd7 100644 --- a/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/litellm_model_name.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; import { getPlaceholder, Providers } from "../provider_info_helpers"; import { MountedFormHost } from "../../../tests/mounted-form-host"; @@ -6,7 +6,7 @@ import LiteLLMModelNameField from "./litellm_model_name"; describe("LitellmModelNameField", () => { it("should render", () => { - const { getByText } = render( + render( { /> , ); - expect(getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); + expect(screen.getByText("LiteLLM Model Name(s)")).toBeInTheDocument(); }); it("should show Azure placeholder as 'my-deployment'", () => { - const { getByPlaceholderText, queryByPlaceholderText } = render( + render( , ); - expect(getByPlaceholderText("my-deployment")).toBeInTheDocument(); - expect(queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("my-deployment")).toBeInTheDocument(); + expect(screen.queryByPlaceholderText("gpt-3.5-turbo")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx b/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx index df4dbff4111..daaacd7f0e2 100644 --- a/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_pass_through.integration.test.tsx @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen, waitFor } from "../../tests/test-utils"; +import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "../../tests/test-utils"; import AddPassThroughEndpoint from "./add_pass_through"; const createPassThroughEndpoint = vi.fn(); @@ -153,8 +153,7 @@ describe("add_pass_through submit payload", () => { await openModal(user); await fillRequiredFields(user); - await user.click(screen.getByLabelText(/HTTP Methods/)); - await user.click(await screen.findByTitle("POST")); + await chooseSelectOption(user, screen.getByLabelText(/HTTP Methods/), "POST"); await submit(user); diff --git a/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx b/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx index 98e876801ac..cff33546fbc 100644 --- a/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx +++ b/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx @@ -5,6 +5,7 @@ import React, { useState, useEffect } from "react"; import { alertingSettingsCall, updateConfigFieldSetting } from "../networking"; import DynamicForm from "./dynamic_form"; +import { extractProxyErrorMessage } from "@/lib/http/client"; import { toast } from "@/lib/toast"; interface alertingSettingsItem { field_name: string; @@ -43,7 +44,7 @@ const AlertingSettings: React.FC = ({ accessToken, premiu setAlertingSettings(updatedSettings); }; - const handleSubmit = (formValues: Record) => { + const handleSubmit = async (formValues: Record) => { if (!accessToken) { return; } @@ -64,18 +65,18 @@ const AlertingSettings: React.FC = ({ accessToken, premiu const mergedFormValues = { ...formValues, ...initialFormValues }; const { slack_alerting, ...alertingArgs } = mergedFormValues; try { - updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs); + await updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs); if (typeof slack_alerting === "boolean") { if (slack_alerting == true) { - updateConfigFieldSetting(accessToken, "alerting", ["slack"]); + await updateConfigFieldSetting(accessToken, "alerting", ["slack"]); } else { - updateConfigFieldSetting(accessToken, "alerting", []); + await updateConfigFieldSetting(accessToken, "alerting", []); } } // update value in state toast.success("Wait 10s for proxy to update."); } catch (error) { - // do something + toast.error(extractProxyErrorMessage(error)); } }; diff --git a/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx b/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx index f9b9765f59e..7764052173c 100644 --- a/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx @@ -38,6 +38,14 @@ const SETTINGS: Setting[] = [ stored_in_db: null, premium_field: false, }, + { + field_name: "daily_spend_per_user_threshold", + field_description: "Daily spend threshold per user", + field_type: "Float", + field_value: 5.5, + stored_in_db: true, + premium_field: false, + }, ]; const renderForm = ( @@ -170,6 +178,19 @@ describe("DynamicForm change notifications", () => { expect(handleInputChange).toHaveBeenCalledWith("daily_report_frequency", 128); }); + it("renders a Float field as a decimal-friendly number input and reports changes as numbers", async () => { + const user = userEvent.setup(); + const { handleInputChange } = renderForm(); + + const input = screen.getByDisplayValue("5.5"); + expect(input).toHaveAttribute("type", "number"); + expect(input).toHaveAttribute("step", "any"); + + await user.type(input, "1"); + + expect(handleInputChange).toHaveBeenCalledWith("daily_spend_per_user_threshold", 5.51); + }); + it("reports a reset with the field name and its row index", async () => { const user = userEvent.setup(); const { handleResetField } = renderForm(); @@ -216,7 +237,7 @@ describe("DynamicForm presentation", () => { expect(screen.getByText("daily_report_frequency")).toBeInTheDocument(); expect(screen.getByText("How often the report runs")).toBeInTheDocument(); - expect(screen.getByText("In DB")).toBeInTheDocument(); + expect(screen.getAllByText("In DB")).toHaveLength(2); expect(screen.getByText("In Config")).toBeInTheDocument(); expect(screen.getByText("Not Set")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx b/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx index be5e76ab0df..42aa58ca0ea 100644 --- a/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx +++ b/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx @@ -63,11 +63,11 @@ const DynamicForm: React.FC = ({ }; const renderControl = (setting: AlertingSetting) => { - if (setting.field_type === "Integer") { + if (setting.field_type === "Integer" || setting.field_type === "Float") { return ( handleNumericChange(setting, event.target.value)} /> diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx index 24685f9c129..e1faf8be1df 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.test.tsx @@ -26,8 +26,8 @@ const openUploadStep = async () => { describe("BulkCreateUsersButton", () => { it("should render", () => { - const { getByText } = render(); - expect(getByText("+ Bulk Invite Users")).toBeInTheDocument(); + render(); + expect(screen.getByText("+ Bulk Invite Users")).toBeInTheDocument(); }); it("parses a CSV chosen through the file input", async () => { diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 25b35c34861..7cbbb08b623 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; import { ArrowLeft, Check, Copy, Link2 } from "lucide-react"; +import { cn } from "@/lib/cva.config"; import { buildMarketplaceSettingsSnippet, formatInstallCommand } from "./helpers"; import { Plugin } from "./types"; @@ -50,48 +51,37 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { ]; return ( -
+
{/* Back link */}
Skills
{/* Header */} -
-

{skill.name}

+
+

{skill.name}

{skill.description && ( -

{skill.description}

+

{skill.description}

)}
{/* Tab bar */} -
-
+
+
{tabs.map((tab) => (
setActiveTab(tab.key)} - style={{ - padding: "12px 20px", - fontSize: 14, - color: activeTab === tab.key ? "#1a73e8" : "#5f6368", - borderBottom: activeTab === tab.key ? "3px solid #1a73e8" : "3px solid transparent", - cursor: "pointer", - fontWeight: activeTab === tab.key ? 500 : 400, - marginBottom: -1, - }} + className={cn( + "-mb-px cursor-pointer border-b-[3px] px-5 py-3 text-sm", + activeTab === tab.key + ? "border-info font-medium text-info" + : "border-transparent font-normal text-muted-foreground", + )} > {tab.label}
@@ -101,27 +91,23 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Overview tab */} {activeTab === "overview" && ( -
+
{/* Left column */} -
-

Skill Details

-

Metadata registered with this skill

-
MetricValue
MetricValue
{row.metric}{row.value}
{row.metric}{row.value}
+
+

Skill Details

+

Metadata registered with this skill

+
- - - + + + {detailRows.map((row, i) => ( - - - + + + ))} @@ -129,38 +115,27 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Right sidebar */} -
-
-
Status
+
+
+
Status
{skill.enabled ? "Public" : "Draft"}
{sourceUrl && ( -
-
Source
+
+
Source
{sourceUrl.replace("https://", "")} @@ -169,20 +144,13 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { )} {skill.keywords && skill.keywords.length > 0 && ( -
-
Tags
-
+
+
Tags
+
{skill.keywords.map((kw) => ( {kw} @@ -192,10 +160,8 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { )}
-
Skill ID
-
- {skill.id} -
+
Skill ID
+
{skill.id}
@@ -203,93 +169,43 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* How to Use tab */} {activeTab === "usage" && ( -
-

Using this skill

-

+

+

Using this skill

+

Once your proxy is set as a marketplace, enable this skill in Claude Code with one command:

{/* Install command */} -
-
- Run in Claude Code +
+
+ Run in Claude Code
- 
-              {installCommand}
-
+ 
{installCommand}
{/* Shown when the marketplace catalog is stale and the plugin isn't found yet */} -
-

+

+

If you see "Plugin {skill.name} not found in marketplace", update the catalog first:

- 
+            
               /plugin marketplace update litellm
-

+

Don't have the marketplace configured yet?{" "} - setActiveTab("setup")} style={{ color: "#1a73e8", cursor: "pointer" }}> + setActiveTab("setup")} className="cursor-pointer text-info"> See one-time setup →

@@ -298,126 +214,56 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { {/* Setup tab (linked from usage) */} {activeTab === "setup" && ( -
-

- One-time marketplace setup -

+
+

One-time marketplace setup

{/* Option 1: single command — fastest path for most users */} -

+

Run this command in Claude Code to register the marketplace:

-
-
- Run in Claude Code +
+
+ Run in Claude Code
- 
+            
               {`/plugin marketplace add ${typeof window !== "undefined" ? window.location.origin : ""}/claude-code/marketplace.json`}
{/* Option 2: settings.json — for persistent config or managed deployments. extraKnownMarketplaces requires source to be a nested object, not a flat string. */} -

- Or add this to{" "} - - ~/.claude/settings.json - {" "} +

+ Or add this to ~/.claude/settings.json{" "} for a persistent configuration:

-
-
- ~/.claude/settings.json +
+
+ ~/.claude/settings.json
- 
-              {settingsSnippet}
-
+ 
{settingsSnippet}
)} diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx index 6e627e3b45c..df21444cbad 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.test.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { Controller, useForm } from "react-hook-form"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; -import { fireEvent, renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import KeyLifecycleSettings from "./KeyLifecycleSettings"; const CREATE_PLACEHOLDER = "e.g., 30d or leave empty to never expire"; @@ -167,8 +167,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("90 days")); + await chooseSelectOption(user, screen.getByRole("combobox"), "90 days"); await waitFor(() => expect(screen.getAllByTitle("90 days").some(isRenderedSelection)).toBe(true)); expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("90d"); @@ -181,8 +180,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); expect(await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).toBeInTheDocument(); expect(screen.getByText("Supported formats: seconds (s), minutes (m), hours (h), days (d)")).toBeInTheDocument(); @@ -196,8 +194,7 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); fireEvent.change(customInput, { target: { value: "14d" } }); @@ -213,14 +210,12 @@ describe("KeyLifecycleSettings", () => { await user.click(screen.getByRole("switch")); expect(await screen.findByText("Rotation Interval")).toBeInTheDocument(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Custom interval")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Custom interval"); const customInput = await screen.findByPlaceholderText("e.g., 1s, 5m, 2h, 14d"); fireEvent.change(customInput, { target: { value: "14d" } }); await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("14d")); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("7 days")); + await chooseSelectOption(user, screen.getByRole("combobox"), "7 days"); await waitFor(() => expect(screen.getByTestId("rotation-interval-value")).toHaveTextContent("7d")); expect(screen.queryByPlaceholderText("e.g., 1s, 5m, 2h, 14d")).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx index bd7b56f1817..7a5db4667e9 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.test.tsx @@ -2,6 +2,7 @@ import { act, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import ModelSelector from "./ModelSelector"; +import { chooseSelectOption } from "../../../tests/test-utils"; vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([]), @@ -9,8 +10,7 @@ vi.mock("@/components/llm_calls/fetch_models", () => ({ const openCustomModelInput = async () => { const user = userEvent.setup(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Enter custom model")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Enter custom model"); return screen.getByPlaceholderText("Enter custom model name"); }; diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx new file mode 100644 index 00000000000..01b8a12d99c --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx @@ -0,0 +1,48 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { chooseSelectOption } from "../../../tests/test-utils"; +import type { Team } from "../key_team_helpers/key_list"; +import TeamDropdown from "./team_dropdown"; + +const TEAMS = [ + { team_id: "team-1", team_alias: "Alpha Team" }, + { team_id: "team-2", team_alias: "Beta Team" }, +] as unknown as Team[]; + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useInfiniteTeams: () => ({ + data: { pages: [{ teams: TEAMS }] }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }), +})); + +describe("TeamDropdown", () => { + it("emits the picked team's id and full object", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await chooseSelectOption(user, screen.getByRole("combobox"), /^Beta Team/); + + expect(onChange).toHaveBeenCalledWith("team-2"); + expect(onTeamSelect).toHaveBeenCalledWith(TEAMS[1]); + }); + + it("emits null, never the empty string, when the selection is cleared", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: "Clear" })); + + expect(onChange).toHaveBeenCalledWith(null); + expect(onTeamSelect).toHaveBeenCalledWith(null); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 35121f41598..7d385c2a3f7 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -5,7 +5,7 @@ import { Team } from "../key_team_helpers/key_list"; interface TeamDropdownProps { value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; /** Callback with the full Team object (or null on clear). */ onTeamSelect?: (team: Team | null) => void; disabled?: boolean; @@ -47,7 +47,7 @@ const TeamDropdown: React.FC = ({ }, [data]); const handleChange = (teamId: string) => { - onChange?.(teamId); + onChange?.(teamId || null); if (onTeamSelect) { onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index d8c2987ead5..3e63cbd3b32 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -257,6 +257,33 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => { }); }); +describe("buildUpdatedComplexityRouterConfig classification mode", () => { + it("round-trips a stored user_turn through hydrate then save", () => { + const stored = { ...STORED, classification_mode: "user_turn" }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + + expect(hydrated.classification_mode).toBe("user_turn"); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classification_mode).toBe("user_turn"); + }); + + it("round-trips an explicitly stored every_request, so an untouched save leaves it as written", () => { + const stored = { ...STORED, classification_mode: "every_request" }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + + expect(hydrated.classification_mode).toBe("every_request"); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classification_mode).toBe("every_request"); + }); + + it("rewrites a stored user_turn to every_request once the operator picks the default back", () => { + const stored = { ...STORED, classification_mode: "user_turn" }; + const result = buildUpdatedComplexityRouterConfig(stored, { + ...FORM_VALUE, + classification_mode: "every_request", + }); + expect(result.classification_mode).toBe("every_request"); + }); +}); + describe("buildUpdatedComplexityRouterConfig deployment affinity", () => { it("writes deployment_affinity=false when the toggle is off", () => { const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: false }); @@ -476,6 +503,7 @@ describe("managed keys survive an untouched open-and-save", () => { classifier_context_budget_chars: 4000, classifier_context_include_assistant_turns: true, classifier_fallback: "default_model", + classification_mode: "user_turn", session_affinity: true, deployment_affinity: false, adaptive: true, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index 027e01a9351..481ba2b6b00 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -1,4 +1,4 @@ -import { buildUpdatedComplexityRouterConfig } from "./edit_auto_router_modal"; +import { buildUpdatedComplexityRouterConfig, hydrateComplexityRouterConfig } from "./edit_auto_router_modal"; const storedConfigValue = { tiers: { @@ -47,8 +47,10 @@ const expectedClassifiedTierConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, adaptive_eligible: "classified_tier", @@ -68,8 +70,10 @@ const expectedAdaptiveDisabledConfig = { semantic_keyword_matching: true, embedding_model: "voyage-4-large", match_threshold: 0.65, + classification_mode: "every_request", session_affinity: false, deployment_affinity: true, + modality_routing: false, }; describe("buildUpdatedComplexityRouterConfig", () => { @@ -85,6 +89,26 @@ describe("buildUpdatedComplexityRouterConfig", () => { expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); }); + it("hydrates a stored modality_routing into form state and defaults absent to off", () => { + expect(hydrateComplexityRouterConfig({ ...storedConfig, modality_routing: true }, null).modality_routing).toBe( + true, + ); + expect(hydrateComplexityRouterConfig(storedConfig, null).modality_routing).toBe(false); + }); + + it("round-trips modality_routing explicitly in both directions", () => { + const enabled = buildUpdatedComplexityRouterConfig(storedConfig, { + ...classifiedTierValue, + modality_routing: true, + }); + expect(enabled.modality_routing).toBe(true); + const disabled = buildUpdatedComplexityRouterConfig( + { ...storedConfig, modality_routing: true }, + { ...classifiedTierValue, modality_routing: false }, + ); + expect(disabled.modality_routing).toBe(false); + }); + it("includes return_raw_model_name only when enabled", () => { const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, { ...classifiedTierValue, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index a4921fcfcb5..96e8549eac8 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -364,7 +364,7 @@ describe("EditAutoRouterModal assistant turns", () => { }); }); -describe("EditAutoRouterModal session affinity", () => { +describe("EditAutoRouterModal classification frequency", () => { beforeEach(() => { modelPatchUpdateCall.mockClear(); }); @@ -381,15 +381,15 @@ describe("EditAutoRouterModal session affinity", () => { />, ); - // A stored config with no session_affinity key now runs with affinity OFF, because the backend - // field defaults to False. The toggle has to render what the router actually does, and an - // untouched save must not flip it. - it("shows a stored config with no session_affinity key as off", async () => { + // A stored config with neither key now runs with affinity OFF, because both backend fields + // default that way. The picker has to render what the router actually does, and an untouched + // save must not flip it. + it("shows a stored config with neither key as every request", async () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every request/ })).toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -397,12 +397,12 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(false); }); - it("shows a stored session_affinity=true as on and preserves it through an untouched save", async () => { + it("shows a stored session_affinity=true as once per session and preserves it through an untouched save", async () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Affinity")); - expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).toBeChecked(); + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Once per session/ })).toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -410,12 +410,12 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(true); }); - it("persists turning session affinity on", async () => { + it("persists picking once per session", async () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Once per session/ })); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -423,18 +423,71 @@ describe("EditAutoRouterModal session affinity", () => { expect(savedConfig().session_affinity).toBe(true); }); - it("persists turning session affinity back off", async () => { + it("persists picking every request back over a stored session pin", async () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Affinity")); - await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every request/ })); await user.click(screen.getByRole("button", { name: /save changes/i })); await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); expect(savedConfig().session_affinity).toBe(false); }); + + it("clears a stored session pin when the operator moves to every new user message", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().session_affinity).toBe(false); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("shows a stored user_turn as selected and preserves it through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, classification_mode: "user_turn" }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + expect(await screen.findByRole("radio", { name: /Every new user message/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("persists switching a stored config to every new user message", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every new user message/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("user_turn"); + }); + + it("rewrites the stored mode to every_request when the operator picks it back", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, classification_mode: "user_turn" }); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("radio", { name: /Every request/ })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classification_mode).toBe("every_request"); + }); }); describe("EditAutoRouterModal deployment affinity", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 4751c2e64b7..e18582f77a0 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -96,11 +96,13 @@ export interface StoredComplexityRouterConfig { classifier_context_budget_chars?: unknown; classifier_context_include_assistant_turns?: unknown; classifier_fallback?: unknown; + classification_mode?: unknown; tier_boundaries?: unknown; token_thresholds?: unknown; dimension_weights?: unknown; reasoning_override_min_score?: unknown; session_affinity?: unknown; + modality_routing?: unknown; deployment_affinity?: unknown; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; @@ -165,12 +167,17 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + classification_mode: + parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" + ? parsedConfig.classification_mode + : undefined, tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), session_affinity: typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, deployment_affinity: typeof parsedConfig.deployment_affinity === "boolean" ? parsedConfig.deployment_affinity @@ -207,7 +214,9 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "classification_mode", "session_affinity", + "modality_routing", "deployment_affinity", "adaptive", "adaptive_weights", @@ -294,6 +303,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, classifierLlmConfig: value.classifier_llm_config, @@ -302,6 +312,7 @@ export const buildUpdatedComplexityRouterConfig = ( classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, classifierFallback: value.classifier_fallback, sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + modalityRouting: value.modality_routing ?? false, deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords: customTechnicalKeywords ?? [], keywordTierRules: keywordMatching?.keywordTierRules ?? [], diff --git a/ui/litellm-dashboard/src/components/model_filters.tsx b/ui/litellm-dashboard/src/components/model_filters.tsx index 96041905d97..0ce82d6d050 100644 --- a/ui/litellm-dashboard/src/components/model_filters.tsx +++ b/ui/litellm-dashboard/src/components/model_filters.tsx @@ -1,5 +1,6 @@ import React, { useState, useEffect, useMemo, useRef } from "react"; import { Card } from "@/components/ui/card"; +import { matchesSearchTerm } from "@/utils/searchUtils"; interface ModelGroupInfo { model_group: string; @@ -76,7 +77,7 @@ const ModelFilters: React.FC = ({ const filteredData = useMemo(() => { return ( modelHubData?.filter((model) => { - const matchesSearch = model.model_group.toLowerCase().includes(searchTerm.toLowerCase()); + const matchesSearch = matchesSearchTerm(searchTerm, [model.model_group]); const matchesProvider = selectedProvider === "" || model.providers.includes(selectedProvider); const matchesMode = selectedMode === "" || model.mode === selectedMode; diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 742d2593ade..768183907db 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -87,6 +87,7 @@ describe("ModelInfoView", () => { accessToken: "test-token", userID: "123", userRole: "Admin", + isViewOnly: false, onModelUpdate: vi.fn(), modelAccessGroups: ["group1", "group2"], }; @@ -328,6 +329,16 @@ describe("ModelInfoView", () => { }); }); + // A proxy_admin_viewer session reads "Admin" through effectiveSessionRole, but the update + // and delete endpoints 403 it, so the write buttons must not be offered. + it("should disable delete and update buttons for a view-only admin session", async () => { + render(, { wrapper }); + await waitFor(() => { + expect(screen.getByTestId("delete-model-button")).toBeDisabled(); + }); + expect(screen.getByTestId("update-api-key-button")).toBeDisabled(); + }); + it("should disable delete button when model is not a DB model", async () => { const nonDbModelData = { ...defaultModelData, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index f21e98e084c..35afcdb2985 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -53,6 +53,7 @@ interface ModelInfoViewProps { accessToken: string | null; userID: string | null; userRole: string | null; + isViewOnly: boolean; onModelUpdate?: (updatedModel: any) => void; modelAccessGroups: string[] | null; } @@ -117,6 +118,7 @@ export default function ModelInfoView({ accessToken, userID, userRole, + isViewOnly, onModelUpdate, modelAccessGroups, }: ModelInfoViewProps) { @@ -167,7 +169,7 @@ export default function ModelInfoView({ // Keep modelData variable name for backwards compatibility const modelData = transformedModelData; - const canEditModel = canModifyModel({ userRole, userID }, teams ?? null, { + const canEditModel = canModifyModel({ userRole, userID, isViewOnly }, teams ?? null, { teamId: modelData?.model_info?.team_id, isDbModel: modelData?.model_info?.db_model === true, }); diff --git a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx index bad92555bd5..3fc1d41dfdd 100644 --- a/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx +++ b/ui/litellm-dashboard/src/components/molecules/cost_optimization_feedback_banner.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { beforeEach, describe, expect, it } from "vitest"; import CostOptimizationFeedbackBanner from "./cost_optimization_feedback_banner"; @@ -10,24 +10,24 @@ describe("CostOptimizationFeedbackBanner", () => { }); it("renders with a link to the feedback discussion", () => { - const { getByText } = render(); - const link = getByText("Share Feedback").closest("a"); + render(); + const link = screen.getByText("Share Feedback").closest("a"); expect(link).toHaveAttribute("href", "https://github.com/BerriAI/litellm/discussions/32172"); }); it("hides itself and persists the dismissal when the dismiss button is clicked", () => { - const { getByText, queryByText, getByLabelText } = render(); - expect(getByText("Help shape cost optimization")).toBeInTheDocument(); + render(); + expect(screen.getByText("Help shape cost optimization")).toBeInTheDocument(); - fireEvent.click(getByLabelText("Dismiss banner")); + fireEvent.click(screen.getByLabelText("Dismiss banner")); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); expect(localStorage.getItem(STORAGE_KEY)).toBe("true"); }); it("stays dismissed on remount once persisted", () => { localStorage.setItem(STORAGE_KEY, "true"); - const { queryByText } = render(); - expect(queryByText("Help shape cost optimization")).not.toBeInTheDocument(); + render(); + expect(screen.queryByText("Help shape cost optimization")).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index b8d8e3ba9c8..700d19eb13c 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -108,7 +108,7 @@ beforeEach(() => { test("renders organization view after loading data", async () => { mockUseOrganization.mockReturnValue({ data: mockOrg, isLoading: false } as any); - const { findAllByText } = renderWithProviders( + renderWithProviders( {}} @@ -120,7 +120,7 @@ test("renders organization view after loading data", async () => { />, ); - const [orgName] = await findAllByText("Acme Corp"); + const [orgName] = await screen.findAllByText("Acme Corp"); expect(orgName).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 5baa8138960..d01a6a34cbe 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -150,6 +150,8 @@ export enum Providers { PETALS = "Petals", PG_VECTOR = "Pg Vector", PREDIBASE = "Predibase", + Qwen_AI_Platform = "Qwen AI Platform", + QwenCloud = "QwenCloud", RECRAFT = "Recraft", REPLICATE = "Replicate", RunwayML = "RunwayML", @@ -262,6 +264,8 @@ export const provider_map: Record = { PETALS: "petals", PG_VECTOR: "pg_vector", PREDIBASE: "predibase", + Qwen_AI_Platform: "qwen_ai_platform", + QwenCloud: "qwencloud", RECRAFT: "recraft", REPLICATE: "replicate", RunwayML: "runwayml", @@ -357,6 +361,8 @@ export const providerLogoMap: Partial> = { [Providers.Openrouter]: openrouterLogo.src, [Providers.Oracle]: oracleLogo.src, [Providers.Perplexity]: perplexityAiLogo.src, + [Providers.Qwen_AI_Platform]: qwenLogo.src, + [Providers.QwenCloud]: qwenLogo.src, [Providers.RECRAFT]: recraftLogo.src, [Providers.REPLICATE]: replicateLogo.src, [Providers.RunwayML]: runwayLogo.src, diff --git a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx index 875f89b5adc..fec46e98077 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.test.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.test.tsx @@ -134,6 +134,27 @@ describe("PublicModelHub", () => { expect(within(gpt35Row as HTMLElement).getByText("Unknown")).toBeInTheDocument(); }); }); + it("shows no models when the search has no matches (LIT-5230 regression)", async () => { + const networkingModule = await import("./networking"); + vi.mocked(networkingModule.modelHubPublicModelsCall).mockResolvedValue([ + { model_group: "gpt-4", providers: ["openai"], mode: "chat" }, + { model_group: "claude-3", providers: ["anthropic"], mode: "chat" }, + ]); + + render(); + expect(await screen.findByText("gpt-4")).toBeInTheDocument(); + + fireEvent.change(screen.getByPlaceholderText("Search model names... (smart search enabled)"), { + target: { value: "zzzz" }, + }); + + await waitFor(() => { + expect(screen.queryByText("gpt-4")).not.toBeInTheDocument(); + expect(screen.queryByText("claude-3")).not.toBeInTheDocument(); + expect(screen.getByText("No matching models")).toBeInTheDocument(); + }); + }); + it("handles non-array response gracefully (regression test for e.filter crash)", async () => { const networkingModule = await import("./networking"); // Mock the API to return an object (like an error response) instead of an array diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 171b7992325..f6364b5d9d1 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -47,6 +47,7 @@ import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets"; import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; import { MessageType } from "@/components/chat_ui/types"; import { getProviderLogoAndName } from "./provider_info_helpers"; +import { filterBySearchTerm, rankBySearchRelevance } from "@/utils/searchUtils"; interface PublicModelHubProps { accessToken?: string | null; @@ -236,52 +237,11 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const filteredData = useMemo(() => { if (!modelHubData || !Array.isArray(modelHubData)) return []; - let searchResults = modelHubData; - - // Apply search if there's a search term - if (searchTerm.trim()) { - const lowercaseSearch = searchTerm.toLowerCase(); - const searchWords = lowercaseSearch.split(/\s+/); - - // First, try flexible matching that handles different separators - const exactMatches = modelHubData.filter((model) => { - const modelName = model.model_group.toLowerCase(); - - // Check if it contains the exact search term - if (modelName.includes(lowercaseSearch)) { - return true; - } - - // Check if it contains all search words (handles spaces vs slashes/dashes) - return searchWords.every((word) => modelName.includes(word)); - }); - - // If we have exact matches, rank them by relevance - if (exactMatches.length > 0) { - searchResults = exactMatches.sort((a, b) => { - const aName = a.model_group.toLowerCase(); - const bName = b.model_group.toLowerCase(); - - // Calculate relevance scores - const aExactMatch = aName === lowercaseSearch ? 1000 : 0; - const bExactMatch = bName === lowercaseSearch ? 1000 : 0; - - const aStartsWith = aName.startsWith(lowercaseSearch) ? 100 : 0; - const bStartsWith = bName.startsWith(lowercaseSearch) ? 100 : 0; - - const aContainsWords = lowercaseSearch.split(/\s+/).every((word) => aName.includes(word)) ? 50 : 0; - const bContainsWords = lowercaseSearch.split(/\s+/).every((word) => bName.includes(word)) ? 50 : 0; - - const aLength = aName.length; - const bLength = bName.length; - - const aScore = aExactMatch + aStartsWith + aContainsWords + (1000 - aLength); - const bScore = bExactMatch + bStartsWith + bContainsWords + (1000 - bLength); - - return bScore - aScore; // Higher score first - }); - } - } + const searchResults = rankBySearchRelevance( + filterBySearchTerm(modelHubData, searchTerm, (model) => [model.model_group]), + searchTerm, + (model) => model.model_group, + ); // Apply other filters return searchResults.filter((model) => { @@ -310,43 +270,11 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const filteredAgentData = useMemo(() => { if (!agentHubData || !Array.isArray(agentHubData)) return []; - let searchResults = agentHubData; - - // Apply search if there's a search term - if (agentSearchTerm.trim()) { - const lowercaseSearch = agentSearchTerm.toLowerCase(); - const searchWords = lowercaseSearch.split(/\s+/); - - searchResults = agentHubData.filter((agent) => { - const agentName = agent.name.toLowerCase(); - const agentDescription = agent.description.toLowerCase(); - - // Check if it contains the exact search term - if (agentName.includes(lowercaseSearch) || agentDescription.includes(lowercaseSearch)) { - return true; - } - - // Check if it contains all search words - return searchWords.every((word) => agentName.includes(word) || agentDescription.includes(word)); - }); - - // Sort by relevance - searchResults = searchResults.sort((a, b) => { - const aName = a.name.toLowerCase(); - const bName = b.name.toLowerCase(); - - const aExactMatch = aName === lowercaseSearch ? 1000 : 0; - const bExactMatch = bName === lowercaseSearch ? 1000 : 0; - - const aStartsWith = aName.startsWith(lowercaseSearch) ? 100 : 0; - const bStartsWith = bName.startsWith(lowercaseSearch) ? 100 : 0; - - const aScore = aExactMatch + aStartsWith + (1000 - aName.length); - const bScore = bExactMatch + bStartsWith + (1000 - bName.length); - - return bScore - aScore; - }); - } + const searchResults = rankBySearchRelevance( + filterBySearchTerm(agentHubData, agentSearchTerm, (agent) => [agent.name, agent.description]), + agentSearchTerm, + (agent) => agent.name, + ); // Apply skill filters return searchResults.filter((agent) => { @@ -361,43 +289,11 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const filteredMcpData = useMemo(() => { if (!mcpHubData || !Array.isArray(mcpHubData)) return []; - let searchResults = mcpHubData; - - // Apply search if there's a search term - if (mcpSearchTerm.trim()) { - const lowercaseSearch = mcpSearchTerm.toLowerCase(); - const searchWords = lowercaseSearch.split(/\s+/); - - searchResults = mcpHubData.filter((server) => { - const serverName = server.server_name.toLowerCase(); - const serverDescription = (server.mcp_info?.description || "").toLowerCase(); - - // Check if it contains the exact search term - if (serverName.includes(lowercaseSearch) || serverDescription.includes(lowercaseSearch)) { - return true; - } - - // Check if it contains all search words - return searchWords.every((word) => serverName.includes(word) || serverDescription.includes(word)); - }); - - // Sort by relevance - searchResults = searchResults.sort((a, b) => { - const aName = a.server_name.toLowerCase(); - const bName = b.server_name.toLowerCase(); - - const aExactMatch = aName === lowercaseSearch ? 1000 : 0; - const bExactMatch = bName === lowercaseSearch ? 1000 : 0; - - const aStartsWith = aName.startsWith(lowercaseSearch) ? 100 : 0; - const bStartsWith = bName.startsWith(lowercaseSearch) ? 100 : 0; - - const aScore = aExactMatch + aStartsWith + (1000 - aName.length); - const bScore = bExactMatch + bStartsWith + (1000 - bName.length); - - return bScore - aScore; - }); - } + const searchResults = rankBySearchRelevance( + filterBySearchTerm(mcpHubData, mcpSearchTerm, (server) => [server.server_name, server.mcp_info?.description]), + mcpSearchTerm, + (server) => server.server_name, + ); // Apply transport filters return searchResults.filter((server) => { diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index b97fef32402..762b23f413e 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -77,21 +77,21 @@ describe("Settings", () => { }); it("should render the logging callbacks tab when access token is provided", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); }); it("should display additional settings tabs", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); - expect(getByText("Alerting Types")).toBeInTheDocument(); - expect(getByText("Alerting Settings")).toBeInTheDocument(); - expect(getByText("Email Alerts")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("Alerting Types")).toBeInTheDocument(); + expect(screen.getByText("Alerting Settings")).toBeInTheDocument(); + expect(screen.getByText("Email Alerts")).toBeInTheDocument(); }); }); @@ -279,13 +279,13 @@ describe("Settings", () => { }); it("should display CloudZero Cost Tracking tab", async () => { - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Active Logging Callbacks")).toBeInTheDocument(); + expect(screen.getByText("Active Logging Callbacks")).toBeInTheDocument(); }); - expect(getByText("CloudZero Cost Tracking")).toBeInTheDocument(); + expect(screen.getByText("CloudZero Cost Tracking")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 72da01d0918..9cb55b6ed5c 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -293,6 +293,8 @@ const Settings: React.FC = ({ accessToken, userRole, userID, llm_too_slow: "LLM Responses Too Slow", llm_requests_hanging: "LLM Requests Hanging", budget_alerts: "Budget Alerts (API Keys, Users)", + user_spend_thresholds: "User Spend Thresholds (Daily/Monthly)", + user_spend_anomalies: "User Spend Anomaly Detection", db_exceptions: "Database Exceptions (Read/Write)", daily_reports: "Weekly/Monthly Spend Reports", outage_alerts: "Outage Alerts", diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 58a0cd94997..f502ea77127 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -1,5 +1,5 @@ import type { ColumnDef, ExpandedState } from "@tanstack/react-table"; -import { render, screen, waitFor } from "@testing-library/react"; +import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; @@ -7,6 +7,7 @@ import { describe, expect, it, vi } from "vitest"; import { DataTable } from "./DataTable"; import { DataTableMultiSortHeader, DataTableSortHeader } from "./DataTableSortHeader"; import { DataTableViewOptions } from "./DataTableViewOptions"; +import { chooseSelectOption } from "../../../../tests/test-utils"; interface Person { id: string; @@ -21,6 +22,12 @@ function person(id: string, name: string, flagged = false): Person { const names = (): (string | null)[] => screen.getAllByTestId("name-cell").map((el) => el.textContent); +const heightClassesOf = (el: HTMLElement | undefined): string[] => + (el?.className ?? "") + .split(/\s+/) + .filter((cls) => cls.startsWith("h-")) + .sort(); + const nameCellColumns: ColumnDef[] = [ { accessorKey: "name", @@ -170,16 +177,13 @@ describe("DataTable sorting", () => { const user = userEvent.setup(); render(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Ascending", "menuitem"); expect(names()).toEqual(["Alice", "Bob", "Charlie"]); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Descending", "menuitem"); expect(names()).toEqual(["Charlie", "Bob", "Alice"]); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Reset")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Reset", "menuitem"); expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); @@ -196,12 +200,10 @@ describe("DataTable sorting", () => { />, ); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Budget descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Budget descending", "menuitem"); expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "max_budget", desc: true }]); - await user.click(screen.getByTestId("sort-trigger-spend")); - await user.click(await screen.findByText("Spend ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Spend ascending", "menuitem"); expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "spend", desc: false }]); }); @@ -229,12 +231,10 @@ describe("DataTable sorting", () => { describe("DataTable layout", () => { it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => { - const { container } = render(); + render(); - const table = container.querySelector("table"); - expect(table).not.toBeNull(); // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow. - expect(table?.style.minWidth).toBe("100%"); + expect(screen.getByRole("table")).toHaveStyle({ minWidth: "100%" }); }); }); @@ -354,17 +354,21 @@ describe("DataTable loading", () => { const { rerender } = render( , ); - const skeletonRow = screen.getAllByTestId("skeleton-row").at(0); - const loadedRowHeight = "h-8"; - expect(skeletonRow?.className).toContain(loadedRowHeight); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); rerender(); - expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight); + const loadedHeight = heightClassesOf(screen.getByRole("row", { name: /Charlie/ })); + + expect(loadedHeight).not.toEqual([]); + expect(skeletonHeight).toEqual(loadedHeight); }); it("does not force the compact height on default-size skeleton rows", () => { - render(); - expect(screen.getAllByTestId("skeleton-row").at(0)?.className).not.toContain("h-8"); + const { rerender } = render(); + const skeletonHeight = heightClassesOf(screen.getAllByRole("row").at(-1)); + + rerender(); + expect(heightClassesOf(screen.getAllByRole("row").at(-1))).not.toEqual(skeletonHeight); }); it("varies skeleton shape and width per column instead of one fixed bar", () => { @@ -420,7 +424,7 @@ describe("DataTable loading", () => { describe("DataTable column visibility", () => { it("hides a column when toggled off in the view-options menu", async () => { const user = userEvent.setup(); - const { container } = render( + render( { />, ); - expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull(); + expect(screen.getByRole("columnheader", { name: "Email" })).toBeInTheDocument(); await user.click(screen.getByTestId("view-options-trigger")); await user.click(await screen.findByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).toBeNull()); + await waitFor(() => expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument()); await user.click(screen.getByTestId("view-option-email")); - await waitFor(() => expect(container.querySelector('th[data-header-id="email"]')).not.toBeNull()); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); }); it("omits columns that opt out of hiding from the menu", async () => { @@ -468,14 +472,10 @@ describe("DataTable column visibility", () => { describe("DataTable pinned columns", () => { it("applies sticky positioning to a pinned column only", () => { - const { container } = render(); + render(); - const pinnedHead = container.querySelector('th[data-header-id="name"]'); - const normalHead = container.querySelector('th[data-header-id="email"]'); - - expect(pinnedHead?.style.position).toBe("sticky"); - expect(pinnedHead?.style.left).toBe("0px"); - expect(normalHead?.style.position).toBe(""); + expect(screen.getByRole("columnheader", { name: "Name" })).toHaveStyle({ position: "sticky", left: "0px" }); + expect(screen.getByRole("columnheader", { name: "Email" })).not.toHaveStyle({ position: "sticky" }); }); }); @@ -570,7 +570,7 @@ describe("DataTable expansion", () => { describe("DataTable row styling and footer", () => { it("applies rowClassName to the matching row only", () => { const data = [person("a", "Alice", true), person("b", "Bob", false)]; - const { container } = render( + render( { />, ); - expect(container.querySelector('tr[data-row-id="a"]')?.className).toContain("flagged-row"); - expect(container.querySelector('tr[data-row-id="b"]')?.className).not.toContain("flagged-row"); + expect(screen.getByRole("row", { name: /Alice/ })).toHaveClass("flagged-row"); + expect(screen.getByRole("row", { name: /Bob/ })).not.toHaveClass("flagged-row"); }); it("renders the footer slot inside a tfoot element", () => { @@ -596,63 +596,57 @@ describe("DataTable row styling and footer", () => { />, ); - expect(screen.getByTestId("footer-row").closest("tfoot")).not.toBeNull(); + const rowGroups = screen.getAllByRole("rowgroup"); + expect(within(rowGroups.at(-1) as HTMLElement).getByText("Total: 3")).toBeInTheDocument(); }); }); describe("DataTable layout", () => { it("exposes resize handles with stable selectors only when resizing is enabled", () => { - const { container, rerender } = render( - , - ); - expect(container.querySelectorAll("[data-resizer][data-header-id]").length).toBe(2); + const { rerender } = render(); + expect(screen.getByTestId("column-resizer-name")).toBeInTheDocument(); + expect(screen.getByTestId("column-resizer-email")).toBeInTheDocument(); rerender(); - expect(container.querySelectorAll("[data-resizer]").length).toBe(0); + expect(screen.queryByTestId("column-resizer-name")).not.toBeInTheDocument(); }); it("makes the header sticky and constrains body height when maxBodyHeight is set", () => { - const { container } = render(); - expect(container.querySelector("thead")?.className).toContain("sticky"); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - expect(scroller).toHaveStyle({ maxHeight: "240px" }); + render(); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky"); + expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" }); }); it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; - const frame = scroller.parentElement as HTMLElement; - const outer = frame.parentElement as HTMLElement; + render(); + const outer = screen.getByTestId("data-table-root"); + const frame = screen.getByTestId("data-table-frame"); + const scroller = screen.getByTestId("data-table-scroller"); // A ceiling, not a stretch: flex-1 here would hold the footer at the bottom on a two-row table. - expect(outer.className).toContain("max-h-full"); - expect(outer.className).not.toContain("flex-1"); - expect(frame.className).not.toContain("flex-1"); - expect(scroller.className).not.toContain("flex-1"); + expect(outer).toHaveClass("max-h-full", "flex-col"); + expect(outer).not.toHaveClass("flex-1"); + expect(frame).toHaveClass("flex-col"); + expect(frame).not.toHaveClass("flex-1"); + expect(scroller).not.toHaveClass("flex-1"); - expect(outer.className).toContain("flex-col"); - expect(frame.className).toContain("flex-col"); - expect(scroller.className).toContain("min-h-0"); - expect(scroller.className).toContain("overflow-auto"); + expect(scroller).toHaveClass("min-h-0", "overflow-auto"); expect(scroller).toHaveStyle({ maxHeight: "" }); // Without this the Table primitive's own overflow container captures the sticky header. - expect(scroller.className).toContain("[&_[data-slot=table-container]]:overflow-visible"); + expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible"); - const thead = container.querySelector("thead") as HTMLElement; - expect(thead.className).toContain("sticky"); // Rows pass under the header, so the semi-transparent row tint alone would let them show through. - expect(thead.className).toContain("bg-background"); + expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background"); }); it("leaves the default layout untouched when neither height mode is set", () => { - const { container } = render(); - const scroller = container.querySelector('[data-slot="table-container"]')?.parentElement as HTMLElement; + render(); + const scroller = screen.getByTestId("data-table-scroller"); - expect(scroller.className).toContain("overflow-x-auto"); - expect(scroller.className).not.toContain("min-h-0"); + expect(scroller).toHaveClass("overflow-x-auto"); + expect(scroller).not.toHaveClass("min-h-0"); expect(scroller).toHaveStyle({ maxHeight: "" }); - expect((scroller.parentElement as HTMLElement).className).not.toContain("flex-col"); - expect(container.querySelector("thead")?.className).not.toContain("sticky"); - expect(container.querySelector("thead")?.className).not.toContain("bg-background"); + expect(screen.getByTestId("data-table-frame")).not.toHaveClass("flex-col"); + expect(screen.getByTestId("data-table-head")).not.toHaveClass("sticky", "bg-background"); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index e8357b37715..60267606951 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -195,8 +195,7 @@ function DataTableHeadCell({ header, size, stickyHeader, enableColumnResi )} {canResize && (
column.resetSize()} @@ -589,15 +588,19 @@ export function DataTable(props: DataTableProps -
+
+
{toolbar !== undefined &&
{toolbar(table)}
}
- + {table.getHeaderGroups().map((headerGroup) => ( {headerGroup.headers.map((header) => ( diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx index 70403ba2efe..aeb6fa44e91 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.test.tsx @@ -13,6 +13,7 @@ import { useState } from "react"; import { describe, expect, it, vi } from "vitest"; import { DataTableSortHeader, type DataTableSortVariant } from "./DataTableSortHeader"; +import { chooseSelectOption } from "../../../../tests/test-utils"; interface Item { name: string; @@ -84,16 +85,13 @@ describe("DataTableSortHeader", () => { const user = userEvent.setup(); render(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Descending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Descending", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="desc"]')).not.toBeNull(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Ascending")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Ascending", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="asc"]')).not.toBeNull(); - await user.click(screen.getByTestId("sort-trigger-name")); - await user.click(await screen.findByText("Reset")); + await chooseSelectOption(user, screen.getByTestId("sort-trigger-name"), "Reset", "menuitem"); expect(screen.getByTestId("sort-trigger-name").querySelector('[data-sort-indicator="none"]')).not.toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx index d248cf311cc..04c8d3e9012 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx @@ -5,6 +5,7 @@ import { describe, expect, it, vi } from "vitest"; import { PaginatedSearchSelect } from "./PaginatedSearchSelect"; import type { SearchSelectOption } from "./SearchSelect"; +import { chooseSelectOption } from "../../../tests/test-utils"; const OPTIONS: SearchSelectOption[] = [ { label: "alias-alpha", value: "alias-alpha" }, @@ -63,8 +64,7 @@ describe("PaginatedSearchSelect", () => { } render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("alias-beta")); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); expect(screen.getByRole("combobox")).toHaveValue("alias-beta"); await new Promise((resolve) => setTimeout(resolve, 400)); @@ -136,8 +136,7 @@ describe("PaginatedSearchSelect", () => { const onValueChange = vi.fn(); renderSelect({ onValueChange }); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("alias-beta")); + await chooseSelectOption(user, screen.getByRole("combobox"), "alias-beta"); expect(onValueChange).toHaveBeenCalledWith("alias-beta"); }); @@ -255,8 +254,7 @@ describe("PaginatedSearchSelect", () => { } render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Beta Team")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Beta Team"); await user.click(screen.getByRole("button", { name: "refetch" })); expect(screen.getByRole("combobox")).toHaveValue("Beta Team"); diff --git a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx index 3698b68155e..5bd48b1aa4c 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginationStatusAlerts.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PaginationStatusAlerts from "./PaginationStatusAlerts"; @@ -6,7 +6,7 @@ import PaginationStatusAlerts from "./PaginationStatusAlerts"; describe("PaginationStatusAlerts", () => { it("shows page progress and wires the Stop button while fetching", () => { const cancel = vi.fn(); - const { getByRole, getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); - fireEvent.click(getByRole("button", { name: "Stop" })); + expect(screen.getByText(/Currently fetching spend data: fetched 7 \/ 42 pages/)).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Stop" })); expect(cancel).toHaveBeenCalledTimes(1); }); it("shows the partial-data notice after a cancel, frozen at the last fetched page", () => { - const { getByText } = render( + render( { />, ); - expect(getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); + expect(screen.getByText("Showing partial spend data (7/42 pages loaded)")).toBeInTheDocument(); }); it("names the subject it is fetching", () => { - const { getByText } = render( + render( { />, ); - expect(getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); + expect(screen.getByText(/Currently fetching agent data: fetched 1 \/ 3 pages/)).toBeInTheDocument(); }); it("renders nothing when idle", () => { diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx index c0b9bc7d8ae..62a510268dd 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import { SearchSelect } from "./SearchSelect"; +import { chooseSelectOption } from "../../../tests/test-utils"; const OPTIONS = [ { label: "Acme Prod", value: "team-1" }, @@ -71,8 +72,7 @@ describe("SearchSelect", () => { const onValueChange = vi.fn(); const user = userEvent.setup(); render(); - await user.click(screen.getByRole("combobox")); - await user.click(await screen.findByText("Growth")); + await chooseSelectOption(user, screen.getByRole("combobox"), "Growth"); expect(onValueChange).toHaveBeenCalledWith("team-2"); }); diff --git a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx index 6e4ab14be33..c81ade526dd 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/area_chart.test.tsx @@ -1,4 +1,4 @@ -import { render } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import React from "react"; import { describe, expect, it } from "vitest"; import { AreaChart } from "./area_chart"; @@ -21,9 +21,9 @@ describe("AreaChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx index a322eeb3ad0..3cd7bd2fd08 100644 --- a/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/charts/bar_chart.test.tsx @@ -21,9 +21,9 @@ describe("BarChart", () => { }); it("renders the No data placeholder instead of a chart when data is empty", () => { - const { container, getByText } = render(); + const { container } = render(); - expect(getByText("No data")).toBeInTheDocument(); + expect(screen.getByText("No data")).toBeInTheDocument(); expect(container.querySelector('[data-slot="chart"]')).toBeNull(); }); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 6592ff8c357..705679fe2e6 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -445,7 +445,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { ); fireEvent.click(screen.getByText("Settings")); - expect(screen.getByText("Budget Reset").parentElement?.textContent).toContain("Every 30d"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Every 30d"); fireEvent.click(screen.getByText("Edit Settings")); (globalThis as any).__TEST_FORM_VALUES = { @@ -456,7 +456,7 @@ describe("KeyInfoView handleKeyUpdate budget_duration", () => { fireEvent.click(screen.getByText("Mock Submit")); await waitFor(() => { - expect(screen.getByText("Budget Reset").parentElement?.textContent).toBe("Budget ResetNever"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index bf8f43b8b5f..97b09e00808 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1,7 +1,7 @@ import { fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "../key_team_helpers/ModelMaxBudgetEditor"; import { @@ -300,7 +300,7 @@ describe("KeyEditView", () => { }); it("should render", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -313,12 +313,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("Save Changes")).toBeInTheDocument(); + expect(screen.getByText("Save Changes")).toBeInTheDocument(); }); }); it("should render tags", async () => { - const { getByText } = renderWithProviders( + renderWithProviders( {}} @@ -331,12 +331,12 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("test-tag")).toBeInTheDocument(); + expect(screen.getByText("test-tag")).toBeInTheDocument(); }); }); it("should not render tags in metadata textarea", async () => { - const { getByLabelText } = renderWithProviders( + renderWithProviders( {}} @@ -348,7 +348,7 @@ describe("KeyEditView", () => { />, ); - const metadataTextarea = getByLabelText("Metadata") as HTMLTextAreaElement; + const metadataTextarea = screen.getByLabelText("Metadata") as HTMLTextAreaElement; await waitFor(() => { expect(metadataTextarea).toHaveValue("{}"); }); @@ -963,10 +963,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - - const weeklyOption = await screen.findByText("weekly"); - await userEvent.click(weeklyOption); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "weekly"); const submitButton = screen.getByRole("button", { name: /save changes/i }); await userEvent.click(submitButton); @@ -1042,8 +1039,7 @@ describe("KeyEditView", () => { ); const resetBudget = await screen.findByLabelText("Reset Budget"); - await userEvent.click(resetBudget); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, resetBudget, "Never resets"); await waitFor(() => { expect(resetBudget).toHaveTextContent("Never resets"); @@ -1074,8 +1070,7 @@ describe("KeyEditView", () => { />, ); - await userEvent.click(await screen.findByLabelText("Reset Budget")); - await userEvent.click(await screen.findByText("Never resets")); + await chooseSelectOption(userEvent, await screen.findByLabelText("Reset Budget"), "Never resets"); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -1946,8 +1941,7 @@ describe("KeyEditView", () => { await userEvent.clear(duration); await userEvent.type(duration, "45d"); - await userEvent.click(screen.getByLabelText(/TPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/TPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); @@ -2103,8 +2097,7 @@ describe("KeyEditView", () => { renderForPayload(onSubmitMock); await screen.findByRole("button", { name: /save changes/i }); - await userEvent.click(screen.getByLabelText(/RPM Rate Limit Type/)); - await userEvent.click(await screen.findByTitle("Guaranteed throughput")); + await chooseSelectOption(userEvent, screen.getByLabelText(/RPM Rate Limit Type/), /^Guaranteed throughput/); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx index f506c0e51d7..a02f80f7800 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx @@ -381,6 +381,6 @@ describe("KeyInfoView budget reset visibility", () => { await waitFor(() => { expect(screen.getByText("Budget Reset")).toBeInTheDocument(); }); - expect(screen.getByText("Budget Reset").parentElement).toHaveTextContent("Never"); + expect(screen.getByTestId("budget-reset-value")).toHaveTextContent("Never"); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0be80c3e173..0dd0dd6d6af 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -895,7 +895,7 @@ export default function KeyInfoView({

Budget Reset

-

+

{currentKeyData.budget_reset_at ? `${currentKeyData.budget_duration ? `Every ${currentKeyData.budget_duration}, next ` : ""}${formatTimestamp(currentKeyData.budget_reset_at)}` : "Never"} diff --git a/ui/litellm-dashboard/src/components/ui/field.test.tsx b/ui/litellm-dashboard/src/components/ui/field.test.tsx index 4bb3d0ec665..5765a637526 100644 --- a/ui/litellm-dashboard/src/components/ui/field.test.tsx +++ b/ui/litellm-dashboard/src/components/ui/field.test.tsx @@ -14,6 +14,7 @@ import { FieldSet, FieldTitle, } from "./field"; +import { ROW_LAYOUT_CLASSES, STRETCH_CHILDREN_CLASS } from "../../../tests/fieldOrientation"; describe("FieldError", () => { it("renders nothing when there are no errors and no children", () => { @@ -85,6 +86,20 @@ describe("Field", () => { expect(screen.getByRole("group")).toHaveAttribute("data-orientation", "horizontal"); }); + + it("stretches every child when vertical, which is what inputs, selects and textareas want", () => { + render(); + + expect(screen.getByRole("group")).toHaveClass("flex-col", STRETCH_CHILDREN_CLASS); + }); + + it("lays children in a row at their own width when horizontal, so a checkbox stays square", () => { + render(); + const field = screen.getByRole("group"); + + expect(field).toHaveClass(...ROW_LAYOUT_CLASSES); + expect(field).not.toHaveClass(STRETCH_CHILDREN_CLASS); + }); }); describe("field primitives forward refs to their DOM node", () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index e0ec66e1ae2..5a35c7ae16b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { useState } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import { ERROR_CODE_OPTIONS } from "./constants"; import { LOG_FILTER_IDS } from "./log_filter_logic"; import { RequestLogsFilters } from "./RequestLogsFilters"; @@ -125,8 +125,7 @@ describe("RequestLogsFilters", () => { const user = userEvent.setup(); const { set } = renderFilters(); - await user.click(await screen.findByPlaceholderText("Search an internal user")); - await user.click(await screen.findByText("alice@example.com")); + await chooseSelectOption(user, await screen.findByPlaceholderText("Search an internal user"), "alice@example.com"); expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "alice@example.com"); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx index 1a6afd7fcfe..8a56312572c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/FormattedToolView.tsx @@ -25,31 +25,15 @@ export function FormattedToolView({ tool }: FormattedToolViewProps) {

{/* Description */} {tool.description && ( -
- - {tool.description} - +
+ {tool.description}
)} {/* Parameters Table */} {parameterRows.length > 0 && (
- - Parameters - + Parameters
- Property - - {skill.name} -
Property{skill.name}
{row.property}{row.value}
{row.property}{row.value}
@@ -82,33 +66,10 @@ export function FormattedToolView({ tool }: FormattedToolViewProps) { {/* If tool was called, show the arguments used */} {tool.called && tool.callData && ( -
- - Called With - -
- 
+        

+          Called With
+          

+            
               {JSON.stringify(tool.callData.arguments, null, 2)}
             

           

diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/JsonToolView.tsx b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/JsonToolView.tsx
index 2a2ceb644dc..d8431e52a51 100644
--- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/JsonToolView.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/JsonToolView.tsx
@@ -20,19 +20,7 @@ export function JsonToolView({ tool }: JsonToolViewProps) {
   };
 
   return (
-    
+    
       {JSON.stringify(toolJson, null, 2)}
     

   );
diff --git a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx
index 112364f5ff3..26b39579859 100644
--- a/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/ToolsSection/ToolItem.tsx
@@ -5,6 +5,7 @@
 import { useState } from "react";
 import { ChevronDown, ChevronRight, Wrench } from "lucide-react";
 import { Badge } from "@/components/ui/badge";
+import { cn } from "@/lib/cva.config";
 import { ParsedTool } from "./types";
 import { ToolExpandedContent } from "./ToolExpandedContent";
 
@@ -16,34 +17,23 @@ export function ToolItem({ tool }: ToolItemProps) {
   const [expanded, setExpanded] = useState(false);
 
   return (
-    

+    

       {/* Header Row - Always Visible */}
       
 setExpanded(!expanded)}
-        style={{
-          display: "flex",
-          alignItems: "center",
-          justifyContent: "space-between",
-          padding: "12px 16px",
-          cursor: "pointer",
-          background: expanded ? "#fafafa" : "#fff",
-          transition: "background 0.2s",
-        }}
+        className={cn(
+          "flex cursor-pointer items-center justify-between gap-3 px-4 py-3 text-card-foreground transition-colors",
+          expanded ? "bg-muted" : "bg-card",
+        )}
       >
-        

+        

           
-          
+          
             {tool.index}. {tool.name}
           
         

 
-        

+        

           {tool.called ? "called" : "not called"}
           {expanded ? (
             
@@ -55,13 +45,7 @@ export function ToolItem({ tool }: ToolItemProps) {
 
       {/* Expanded Content */}
       {expanded && (
-        

+        

           
         

       )}
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
index f14d6279e32..2a8306473b8 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
@@ -129,6 +129,18 @@ describe("autorouter_presets", () => {
     }
   });
 
+  it("carries a preset's modality_routing into the prefilled form state", () => {
+    const preset = getPresetByKey("anthropic_family")!;
+    const withFlag = { ...preset.complexity_router_config, modality_routing: true };
+    const prefill = buildPresetPrefill(withFlag, groupsOnly(getRequiredModelsInPreset(preset)));
+    expect(prefill.complexityRouterConfig.modality_routing).toBe(true);
+    const withoutFlag = buildPresetPrefill(
+      preset.complexity_router_config,
+      groupsOnly(getRequiredModelsInPreset(preset)),
+    );
+    expect(withoutFlag.complexityRouterConfig.modality_routing).toBe(false);
+  });
+
   it("prefills the anthropic preset's effort through to tier_model_params", () => {
     const preset = getPresetByKey("anthropic_family")!;
     const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset)));
@@ -230,6 +242,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: [presetModel], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -273,6 +286,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: ["claude-opus-5"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -287,6 +301,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: ["claude-opus-5"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -320,6 +335,7 @@ describe("autorouter_presets", () => {
     const simpleTierConfig = (presetModel: string) => ({
       tiers: { SIMPLE: [presetModel], MEDIUM: [], COMPLEX: [], REASONING: [] },
       classifier_type: "heuristic" as const,
+      classification_mode: "every_request" as const,
       session_affinity: false,
       deployment_affinity: true,
     });
@@ -563,6 +579,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
         match_threshold: 0,
@@ -578,6 +595,7 @@ describe("autorouter_presets", () => {
         {
           tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
           classifier_type: "heuristic",
+          classification_mode: "every_request",
           session_affinity: false,
           deployment_affinity: true,
           enable_context_window_escalation: false,
@@ -589,11 +607,29 @@ describe("autorouter_presets", () => {
       expect(prefill.complexityRouterConfig.context_window_escalation_buffer).toBe(0.9);
     });
 
+    it("carries a preset's classification_mode and defaults it when the preset omits one", () => {
+      const tiers = { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] };
+      const base = {
+        tiers,
+        classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
+        session_affinity: false,
+        deployment_affinity: true,
+      };
+      const availability = groupsOnly(["gpt-5-nano"]);
+      expect(
+        buildPresetPrefill({ ...base, classification_mode: "user_turn" }, availability).complexityRouterConfig
+          .classification_mode,
+      ).toBe("user_turn");
+      expect(buildPresetPrefill(base, availability).complexityRouterConfig.classification_mode).toBe("every_request");
+    });
+
     it("falls back to the defaults when a preset omits match_threshold and escalation_keywords", () => {
       const prefill = buildPresetPrefill(
         {
           tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
           classifier_type: "heuristic",
+          classification_mode: "every_request" as const,
           session_affinity: false,
           deployment_affinity: true,
         },
@@ -610,6 +646,7 @@ describe("autorouter_presets", () => {
       const base = {
         tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -625,6 +662,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: ["claude-sonnet-4-5"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -639,6 +677,7 @@ describe("autorouter_presets", () => {
           REASONING: [{ model_name: "o3", litellm_params: { reasoning_effort: "high" } }],
         },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -658,6 +697,7 @@ describe("autorouter_presets", () => {
           REASONING: [{ model_name: "claude-sonnet-4-5", litellm_params: { reasoning_effort: "high" } }],
         },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
@@ -680,6 +720,9 @@ describe("autorouter_presets", () => {
           ],
         },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
+        session_affinity: false,
+        deployment_affinity: true,
       };
       const prefill = buildPresetPrefill(config, groupsOnly(["claude-sonnet-4.5"]));
       // temperature survives from the spelling that would otherwise have been overwritten;
@@ -693,6 +736,7 @@ describe("autorouter_presets", () => {
       const config = {
         tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] },
         classifier_type: "heuristic" as const,
+        classification_mode: "every_request" as const,
         session_affinity: false,
         deployment_affinity: true,
       };
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
index f96dd5ddb4c..721bd6f2b2a 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
@@ -6,6 +6,7 @@ import {
   ComplexityRouterConfigValue,
   ClassifierType,
   ClassifierLLMConfig,
+  DEFAULT_CLASSIFICATION_MODE,
   DEFAULT_SESSION_AFFINITY,
   DEFAULT_DEPLOYMENT_AFFINITY,
   usesLlmClassifier,
@@ -288,8 +289,10 @@ export const buildPresetPrefill = (
       classifier_context_budget_chars: config.classifier_context_budget_chars,
       classifier_context_per_turn_chars: config.classifier_context_per_turn_chars,
       classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns,
+      classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE,
       session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
       deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
+      modality_routing: config.modality_routing ?? false,
       adaptive: config.adaptive,
       adaptive_weights: config.adaptive_weights,
       tier_distance_penalty: config.tier_distance_penalty,
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index fb81191cfbf..6f044fec3f3 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -7710,6 +7710,10 @@ export interface paths {
          *         - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
          *         - model_max_budget_usage: dict | None - Current-window spend per model, present only when
          *           the key has per-model budgets
+         *         - budget_limits: list | None - Concurrent budget windows, exactly as stored
+         *         - budget_limits_usage: dict | None - Current-window spend per budget window, e.g.
+         *           {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows
+         *           (read from the same cross-pod spend counter the budget enforcement uses)
          *         - models: list - Model_name's the key is allowed to call
          *         - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits
          *         - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"}
@@ -13285,6 +13289,58 @@ export interface paths {
         patch: operations["patch_user_scim_v2_Users__user_id__patch"];
         trace?: never;
     };
+    "/scim/v2/placeholders": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        /**
+         * List Placeholders
+         * @description List user rows whose id is another account's SSO identity or email.
+         *
+         *     An earlier release provisioned a group member it could not match as a user keyed
+         *     by the raw member value, and that row now shadows the account the value really
+         *     names, so every push of that member is refused. This lists those rows so an
+         *     operator can fold each one into the account it shadows with
+         *     ``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of
+         *     its own or owns virtual keys is left out: someone uses that account.
+         */
+        get: operations["list_placeholders_scim_v2_placeholders_get"];
+        put?: never;
+        post?: never;
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
+    "/scim/v2/placeholders/{user_id}/merge": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        get?: never;
+        put?: never;
+        /**
+         * Merge Placeholder
+         * @description Fold a placeholder user into the one account its id names by SSO identity or email.
+         *
+         *     The account is added to every team the placeholder is on, then the placeholder is
+         *     deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group
+         *     push resolves the member value to the real account. Refused with 409 when the row
+         *     has an SSO identity of its own, owns virtual keys, or names no account or several.
+         */
+        post: operations["merge_placeholder_scim_v2_placeholders__user_id__merge_post"];
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
     "/search": {
         parameters: {
             query?: never;
@@ -22958,7 +23014,7 @@ export interface components {
          * @description Enum for alert types and management event types
          * @enum {string}
          */
-        AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted";
+        AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "user_spend_thresholds" | "user_spend_anomalies" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted";
         /** AllowedVectorStoreIndexItem */
         AllowedVectorStoreIndexItem: {
             /** Index Name */
@@ -25431,9 +25487,15 @@ export interface components {
             database_extra_connection_params?: {
                 [key: string]: unknown;
             } | null;
+            /**
+             * Database Max Idle Connection Lifetime
+             * @description Prisma `max_idle_connection_lifetime` URL param (seconds). A pooled connection idle longer than this is closed and replaced instead of being handed to the next request. Defaults to 60 so connections are recycled before common infra idle timeouts (AWS NLB / RDS Proxy ~350s, many LBs 60-350s) silently drop them and requests fail with `Error { kind: Closed }`. A value pinned on the DATABASE_URL or set via `database_extra_connection_params` takes precedence.
+             * @default 60
+             */
+            database_max_idle_connection_lifetime: number | null;
             /**
              * Database Socket Timeout
-             * @description Prisma `socket_timeout` URL param (seconds). When set, an idle/slow connection that has not produced data within this window is closed. This is the main knob for capping idle DB connections from LiteLLM.
+             * @description Prisma `socket_timeout` URL param (seconds). When set, an in-flight operation that has not produced data within this window is aborted. For capping how long idle pooled connections are kept, see `database_max_idle_connection_lifetime`.
              */
             database_socket_timeout?: number | null;
             /**
@@ -34917,6 +34979,27 @@ export interface components {
             /** Value */
             value?: unknown | null;
         };
+        /**
+         * SCIMPlaceholder
+         * @description A user row keyed by a value that names another account by SSO identity or email.
+         */
+        SCIMPlaceholder: {
+            /** Placeholder User Id */
+            placeholder_user_id: string;
+            /** Resolved User Ids */
+            resolved_user_ids: string[];
+            /** Team Ids */
+            team_ids: string[];
+        };
+        /** SCIMPlaceholderMergeResult */
+        SCIMPlaceholderMergeResult: {
+            /** Merged Into User Id */
+            merged_into_user_id: string;
+            /** Placeholder User Id */
+            placeholder_user_id: string;
+            /** Team Ids */
+            team_ids: string[];
+        };
         /** SCIMServiceProviderConfig */
         SCIMServiceProviderConfig: {
             /** Authenticationschemes */
@@ -35327,8 +35410,17 @@ export interface components {
             last_error?: string | null;
             /** @description Stratified verdicts; detail endpoint only */
             results?: components["schemas"]["ShadowEvalResult"] | null;
-            /** Router Name */
-            router_name: string;
+            /**
+             * Router Name
+             * @description The first router, kept for callers that predate router_names; derived so the
+             *     two fields can never disagree.
+             */
+            readonly router_name: string;
+            /**
+             * Router Names
+             * @description Every auto-router this job runs as a shadow arm. Multi-router jobs sample one slice of traffic and judge every arm against the same real responses
+             */
+            router_names: string[];
             /** Shadow Percentage */
             shadow_percentage: number;
             /**
@@ -35416,6 +35508,12 @@ export interface components {
              * @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked
              */
             by_current_model: components["schemas"]["ShadowEvalSlice"][];
+            /**
+             * By Router
+             * @description One slice per router arm, grouped on the router name. Every arm of a multi-router job is judged against the same real responses over the same sampled requests, so these slices compare routers head-to-head: like-for-like win rates and spends on identical traffic. Verdicts from before arm stamping existed count toward the job's own router
+             * @default []
+             */
+            by_router: components["schemas"]["ShadowEvalSlice"][];
             /** By Tier */
             by_tier: components["schemas"]["ShadowEvalSlice"][];
             /**
@@ -35429,13 +35527,13 @@ export interface components {
             overall_tie_rate_pct: number;
             /**
              * Sampled Real Spend
-             * @description USD the real arm billed across all judged turns, cache-served turns excluded
+             * @description USD the real arm billed across all judged turns, cache-served turns excluded. A judged turn is one (request, router arm) verdict, so a multi-router job counts the real response once per arm it was judged against; per-router comparisons read by_router
              * @default 0
              */
             sampled_real_spend: number;
             /**
              * Sampled Shadow Spend
-             * @description USD the shadow arm billed across the same turns, judge excluded, like for like
+             * @description USD the shadow arms billed across the same turns, judge excluded, like for like
              * @default 0
              */
             sampled_shadow_spend: number;
@@ -35729,15 +35827,21 @@ export interface components {
             judge_model: string;
             /**
              * Max Budget
-             * @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window
+             * @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window. Every router arm draws from the same per-target budget, so a multi-router job reaches it proportionally sooner
              * @default 10
              */
             max_budget: number;
             /**
              * Router Name
-             * @description The auto-router under evaluation, in either direction
+             * @description The auto-router under evaluation, in either direction: the single-router spelling of router_names. Provide exactly one of the two fields
              */
-            router_name: string;
+            router_name?: string | null;
+            /**
+             * Router Names
+             * @description The auto-routers under evaluation, at most 4. Every sampled request runs through every router listed and each arm is judged independently against the same real response, so routers compare head-to-head on identical traffic. More than one router requires direction 'forward'. After validation this field always carries the full deduplicated set, whichever spelling the caller used
+             * @default []
+             */
+            router_names: string[];
             /**
              * Shadow Percentage
              * @description Percentage of each target's requests to duplicate through the router
@@ -55827,6 +55931,70 @@ export interface operations {
             };
         };
     };
+    list_placeholders_scim_v2_placeholders_get: {
+        parameters: {
+            query?: {
+                feature?: string | null;
+            };
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["SCIMPlaceholder"][];
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
+    merge_placeholder_scim_v2_placeholders__user_id__merge_post: {
+        parameters: {
+            query?: {
+                feature?: string | null;
+            };
+            header?: never;
+            path: {
+                user_id: string;
+            };
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["SCIMPlaceholderMergeResult"];
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
     search_search_post: {
         parameters: {
             query?: {
diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts
index 179a9b7933a..afc2ecd210f 100644
--- a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts
+++ b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts
@@ -1,14 +1,16 @@
 import { describe, expect, it } from "vitest";
 
 import { Team } from "@/components/networking";
-import { canModifyModel, modelCreationScope } from "./modelPermissions";
+import { canCreateModels, canModifyModel, modelCreationScope } from "./modelPermissions";
 
 const teamWhere = (userId: string, role: string, teamId = "team-1"): Team[] =>
   [{ team_id: teamId, members_with_roles: [{ user_id: userId, user_email: "t@test.com", role }] }] as unknown as Team[];
 
-const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin" };
-const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" };
-const MEMBER = { userRole: "Internal User", userID: "u-member" };
+const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false };
+const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false };
+const MEMBER = { userRole: "Internal User", userID: "u-member", isViewOnly: false };
+// proxy_admin_viewer sessions: effectiveSessionRole masquerades the role as "Admin".
+const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true };
 
 const noLimits = { disabledForInternalUsers: false };
 
@@ -40,9 +42,24 @@ describe("modelCreationScope", () => {
   // an unscoped create from them 403s. Treating them as admins here is what let a form submit
   // a payload the backend always rejected.
   it("does not treat an org admin as able to create unscoped", () => {
-    const orgAdmin = { userRole: "org_admin", userID: "u-org" };
+    const orgAdmin = { userRole: "org_admin", userID: "u-org", isViewOnly: false };
     expect(modelCreationScope(orgAdmin, { teams: teamWhere("u-org", "admin"), ...noLimits })).toBe("team-required");
   });
+
+  // Server-side, POST /model/new 403s the viewer roles, so the "Admin" the masquerade
+  // reports must not read as a proxy admin here.
+  it("forbids a view-only admin session despite the masqueraded Admin role", () => {
+    expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe("forbidden");
+    expect(canCreateModels(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe(false);
+  });
+
+  // _check_proxy_admin_viewer_access (route_checks.py) 403s /model/new on the session role
+  // alone, before the team-scoped carve-out in ModelManagementAuthChecks can run.
+  it("forbids a view-only admin even when they admin a team", () => {
+    expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: teamWhere("u-viewer", "admin"), ...noLimits })).toBe(
+      "forbidden",
+    );
+  });
 });
 
 describe("canModifyModel", () => {
@@ -80,6 +97,17 @@ describe("canModifyModel", () => {
   });
 
   it("does not treat two absent identities as a match", () => {
-    expect(canModifyModel({ userRole: "Internal User", userID: null }, null, teamRow)).toBe(false);
+    expect(canModifyModel({ userRole: "Internal User", userID: null, isViewOnly: false }, null, teamRow)).toBe(false);
+  });
+
+  // PATCH /model/{id}/update and POST /model/delete 403 the viewer roles like /model/new does.
+  it("refuses a view-only admin session on a DB row", () => {
+    expect(canModifyModel(VIEW_ONLY_ADMIN, null, teamRow)).toBe(false);
+  });
+
+  // The route RBAC blocks /model/update and /model/delete for the viewer role before the
+  // team-scoped carve-out runs, so team-admin membership changes nothing here either.
+  it("refuses a view-only user even when they admin the owning team", () => {
+    expect(canModifyModel(VIEW_ONLY_ADMIN, teamWhere("u-viewer", "admin"), teamRow)).toBe(false);
   });
 });
diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.ts b/ui/litellm-dashboard/src/utils/modelPermissions.ts
index b5914f9d7ea..843d9041026 100644
--- a/ui/litellm-dashboard/src/utils/modelPermissions.ts
+++ b/ui/litellm-dashboard/src/utils/modelPermissions.ts
@@ -3,20 +3,32 @@ import { Team } from "@/components/networking";
 import { isProxyAdminRole, isUserTeamAdminForAnyTeam, isUserTeamAdminForSingleTeam } from "./roles";
 
 /**
- * The dashboard's mirror of ModelManagementAuthChecks in
- * litellm/proxy/management_endpoints/model_management_endpoints.py.
+ * The dashboard's mirror of the two server layers that gate model writes: the role-level
+ * route RBAC (`_check_proxy_admin_viewer_access` in litellm/proxy/auth/route_checks.py),
+ * which 403s /model/new, /model/update, and /model/delete for every view-only session
+ * before the endpoint runs, and ModelManagementAuthChecks in
+ * litellm/proxy/management_endpoints/model_management_endpoints.py behind it.
  *
- * Both questions below are answered there by exactly two inputs: the caller's role, and
- * whether the caller admins the team named in `model_info.team_id`. `created_by` is written
- * at creation and never read by an auth check, so it is deliberately absent here; gating on
- * it hid controls from team admins the API accepts, and showed controls to former team admins
- * the API rejects.
+ * Past that route gate, both questions below are answered by exactly two inputs: the
+ * caller's role, and whether the caller admins the team named in `model_info.team_id`.
+ * `created_by` is written at creation and never read by an auth check, so it is deliberately
+ * absent here; gating on it hid controls from team admins the API accepts, and showed
+ * controls to former team admins the API rejects.
  */
 export interface ModelActor {
   userRole: string | null;
   userID: string | null;
+  /**
+   * From useAuthorized(). A proxy_admin_viewer session masquerades as "Admin" in userRole
+   * (effectiveSessionRole, for read parity), yet every management write 403s it, so the role
+   * alone cannot answer a write question.
+   */
+  isViewOnly: boolean;
 }
 
+const isWritableProxyAdmin = ({ userRole, isViewOnly }: ModelActor): boolean =>
+  !isViewOnly && userRole != null && isProxyAdminRole(userRole);
+
 /** How this actor must scope a deployment they create, or that they may not create one. */
 export type ModelWriteScope = "forbidden" | "unscoped-ok" | "team-required";
 
@@ -33,20 +45,25 @@ const isTeamAdminOf = (teams: Team[] | null, userID: string, teamId: string): bo
 
 /**
  * POST /model/new takes a proxy admin unconditionally, or a team admin whose payload names a
- * team; an unscoped create from anyone else is a 403. Returning the requirement rather than a
- * pair of booleans keeps "may not create" and "may create unscoped" from being confused.
+ * team; an unscoped create from anyone else is a 403. A view-only session is 403d by the
+ * route RBAC on its role alone, so team-admin membership cannot rescue it. Returning the
+ * requirement rather than a pair of booleans keeps "may not create" and "may create
+ * unscoped" from being confused.
  */
 export const modelCreationScope = (
-  { userRole, userID }: ModelActor,
+  actor: ModelActor,
   { teams, disabledForInternalUsers }: ModelCreationLimits,
 ): ModelWriteScope => {
-  if (userRole != null && isProxyAdminRole(userRole)) {
+  if (actor.isViewOnly) {
+    return "forbidden";
+  }
+  if (isWritableProxyAdmin(actor)) {
     return "unscoped-ok";
   }
   if (disabledForInternalUsers) {
     return "forbidden";
   }
-  if (userID != null && isUserTeamAdminForAnyTeam(teams, userID)) {
+  if (actor.userID != null && isUserTeamAdminForAnyTeam(teams, actor.userID)) {
     return "team-required";
   }
   return "forbidden";
@@ -63,18 +80,18 @@ export interface ModelRowOrigin {
 
 /** May this actor edit or delete this specific deployment? */
 export const canModifyModel = (
-  { userRole, userID }: ModelActor,
+  actor: ModelActor,
   teams: Team[] | null,
   { teamId, isDbModel }: ModelRowOrigin,
 ): boolean => {
-  if (!isDbModel) {
+  if (actor.isViewOnly || !isDbModel) {
     return false;
   }
-  if (userRole != null && isProxyAdminRole(userRole)) {
+  if (isWritableProxyAdmin(actor)) {
     return true;
   }
-  if (userID == null || teamId == null) {
+  if (actor.userID == null || teamId == null) {
     return false;
   }
-  return isTeamAdminOf(teams, userID, teamId);
+  return isTeamAdminOf(teams, actor.userID, teamId);
 };
diff --git a/ui/litellm-dashboard/src/utils/searchUtils.test.ts b/ui/litellm-dashboard/src/utils/searchUtils.test.ts
new file mode 100644
index 00000000000..4935e8ab15d
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/searchUtils.test.ts
@@ -0,0 +1,64 @@
+import { describe, expect, it } from "vitest";
+
+import { filterBySearchTerm, matchesSearchTerm, rankBySearchRelevance } from "./searchUtils";
+
+describe("matchesSearchTerm", () => {
+  it("matches everything on an empty or whitespace-only term", () => {
+    expect(matchesSearchTerm("", ["anything"])).toBe(true);
+    expect(matchesSearchTerm("   ", ["anything"])).toBe(true);
+  });
+
+  it("matches a substring of any field, case-insensitively", () => {
+    expect(matchesSearchTerm("BILL", ["Billing Router", "routes invoices"])).toBe(true);
+    expect(matchesSearchTerm("invoice", ["Billing Router", "routes invoices"])).toBe(true);
+  });
+
+  it("matches when every word appears in some field", () => {
+    expect(matchesSearchTerm("router invoices", ["Billing Router", "routes invoices"])).toBe(true);
+    expect(matchesSearchTerm("router refunds", ["Billing Router", "routes invoices"])).toBe(false);
+  });
+
+  it("returns false when nothing matches", () => {
+    expect(matchesSearchTerm("zzzz", ["Billing Router", "routes invoices"])).toBe(false);
+  });
+
+  it("ignores null and undefined fields", () => {
+    expect(matchesSearchTerm("billing", [null, undefined, "Billing Router"])).toBe(true);
+    expect(matchesSearchTerm("billing", [null, undefined])).toBe(false);
+  });
+});
+
+describe("filterBySearchTerm", () => {
+  const agents = [
+    { name: "Billing Router", description: "routes invoices" },
+    { name: "Support Bot", description: "handles tickets" },
+  ];
+
+  it("keeps only items whose fields match", () => {
+    expect(filterBySearchTerm(agents, "tickets", (a) => [a.name, a.description])).toEqual([agents[1]]);
+  });
+
+  it("returns an empty list when nothing matches", () => {
+    expect(filterBySearchTerm(agents, "zzzz", (a) => [a.name, a.description])).toEqual([]);
+  });
+
+  it("returns all items for an empty term", () => {
+    expect(filterBySearchTerm(agents, "", (a) => [a.name, a.description])).toEqual(agents);
+  });
+});
+
+describe("rankBySearchRelevance", () => {
+  it("orders exact match, then prefix match, then shorter names", () => {
+    const items = [{ name: "gpt-4o-mini-transcribe" }, { name: "gpt-4o" }, { name: "chatgpt-4o-latest" }];
+    expect(rankBySearchRelevance(items, "gpt-4o", (m) => m.name).map((m) => m.name)).toEqual([
+      "gpt-4o",
+      "gpt-4o-mini-transcribe",
+      "chatgpt-4o-latest",
+    ]);
+  });
+
+  it("keeps the original order for an empty term", () => {
+    const items = [{ name: "b" }, { name: "a" }];
+    expect(rankBySearchRelevance(items, "", (m) => m.name)).toEqual(items);
+  });
+});
diff --git a/ui/litellm-dashboard/src/utils/searchUtils.ts b/ui/litellm-dashboard/src/utils/searchUtils.ts
new file mode 100644
index 00000000000..b256a5db0c4
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/searchUtils.ts
@@ -0,0 +1,32 @@
+type SearchField = string | null | undefined;
+
+const normalizeTerm = (term: string): string => term.trim().toLowerCase();
+
+export function matchesSearchTerm(term: string, fields: ReadonlyArray): boolean {
+  const needle = normalizeTerm(term);
+  if (needle === "") return true;
+
+  const haystacks = fields.filter((field): field is string => typeof field === "string").map((f) => f.toLowerCase());
+  if (haystacks.some((haystack) => haystack.includes(needle))) return true;
+
+  return needle.split(/\s+/).every((word) => haystacks.some((haystack) => haystack.includes(word)));
+}
+
+export function filterBySearchTerm(
+  items: ReadonlyArray,
+  term: string,
+  fields: (item: T) => ReadonlyArray,
+): T[] {
+  return items.filter((item) => matchesSearchTerm(term, fields(item)));
+}
+
+export function rankBySearchRelevance(items: ReadonlyArray, term: string, name: (item: T) => string): T[] {
+  const needle = normalizeTerm(term);
+  if (needle === "") return [...items];
+
+  const score = (item: T): number => {
+    const candidate = name(item).toLowerCase();
+    return (candidate === needle ? 1000 : 0) + (candidate.startsWith(needle) ? 100 : 0) + (1000 - candidate.length);
+  };
+  return [...items].sort((a, b) => score(b) - score(a));
+}
diff --git a/ui/litellm-dashboard/tests/fieldOrientation.ts b/ui/litellm-dashboard/tests/fieldOrientation.ts
new file mode 100644
index 00000000000..49ba5056d66
--- /dev/null
+++ b/ui/litellm-dashboard/tests/fieldOrientation.ts
@@ -0,0 +1,18 @@
+import { expect } from "vitest";
+
+export const ROW_LAYOUT_CLASSES = ["flex-row", "items-center"] as const;
+export const STRETCH_CHILDREN_CLASS = "*:w-full";
+
+/**
+ * Asserts a control sits beside its label at its own width instead of being stretched across the
+ * field. Reaches for the resolved classes because the defect is purely visual: nothing accessible
+ * distinguishes a square checkbox from a full-width bar.
+ */
+export const expectControlBesideLabel = (control: HTMLElement): void => {
+  const field = control.closest('[data-slot="field"]');
+  if (field === null) throw new Error("control is not rendered inside a form field");
+
+  expect(field).toHaveAttribute("data-orientation", "horizontal");
+  expect(field).toHaveClass(...ROW_LAYOUT_CLASSES);
+  expect(field).not.toHaveClass(STRETCH_CHILDREN_CLASS);
+};
diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx
index 66966201a9c..553726faff0 100644
--- a/ui/litellm-dashboard/tests/test-utils.tsx
+++ b/ui/litellm-dashboard/tests/test-utils.tsx
@@ -44,20 +44,14 @@ const pointerBlocked = (element: HTMLElement): boolean => {
   return false;
 };
 
-/**
- * Opens a Base UI Select and picks an option by its accessible name.
- *
- * The option is in the DOM one render before the popup finishes entering, and until then its
- * positioner still carries `pointer-events: none`, which user-event refuses to click. Waiting on
- * the option text alone is a race that React 19's flush timing loses.
- */
 export const chooseSelectOption = async (
-  user: ReturnType,
+  user: Pick, "click">,
   trigger: HTMLElement,
   optionName: string | RegExp,
+  role: "option" | "menuitem" | "menuitemradio" = "option",
 ) => {
   await user.click(trigger);
-  const option = await screen.findByRole("option", { name: optionName });
+  const option = await screen.findByRole(role, { name: optionName });
   await waitFor(() => expect(pointerBlocked(option)).toBe(false));
   await user.click(option);
 };
diff --git a/uv.lock b/uv.lock
index 7df1a35b28a..27be919eea1 100644
--- a/uv.lock
+++ b/uv.lock
@@ -10,7 +10,7 @@ resolution-markers = [
 ]
 
 [options]
-exclude-newer = "2026-08-26T18:33:25.773031Z"
+exclude-newer = "2026-08-29T17:58:57.633306Z"
 exclude-newer-span = "P3D"
 
 [manifest]
@@ -4266,7 +4266,7 @@ wheels = [
 
 [[package]]
 name = "litellm"
-version = "1.100.0"
+version = "1.101.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiohttp" },
@@ -4669,12 +4669,12 @@ proxy-dev = [
 
 [[package]]
 name = "litellm-enterprise"
-version = "0.1.62"
+version = "0.1.63"
 source = { editable = "enterprise" }
 
 [[package]]
 name = "litellm-proxy-extras"
-version = "0.4.91"
+version = "0.4.92"
 source = { editable = "litellm-proxy-extras" }
 
 [[package]]
diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt
index 7e20081988d..8753d7c3c77 100644
--- a/whitelisted_bedrock_models.txt
+++ b/whitelisted_bedrock_models.txt
@@ -217,3 +217,17 @@ bedrock/us-east-1/zai.glm-5
 bedrock/us-west-2/zai.glm-5
 bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0
 bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0
+bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b
+bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2
+bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b
+bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0
+bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0
+bedrock/us-gov-west-1/anthropic.claude-sonnet-5
+bedrock/us-gov-west-1/anthropic.claude-opus-4-8
+bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b
+bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2
+bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b
+bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0
+bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0
+bedrock/us-gov-east-1/anthropic.claude-sonnet-5
+bedrock/us-gov-east-1/anthropic.claude-opus-4-8