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/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 df52069e71f..da788bf1ce3 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -117,7 +117,7 @@ "limit": 111 }, "reportUnnecessaryComparison": { - "limit": 695 + "limit": 692 }, "reportUnnecessaryContains": { "limit": 5 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/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/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/litellm/constants.py b/litellm/constants.py index 9a50797f517..1bd977dd9a9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,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. @@ -1589,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/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/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index c137164ecdb..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 @@ -1944,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 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/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index e87ac9521ae..372c9bf6b91 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,4 +1,5 @@ import contextvars +import copy import hashlib import os import secrets @@ -39,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() @@ -852,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. diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index e5789965c6e..5e116b7301a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -11,6 +11,7 @@ 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 @@ -30,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, @@ -44,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): @@ -222,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, @@ -241,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", @@ -314,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 @@ -335,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): def _get_response_messages( self, standard_logging_payload: StandardLoggingPayload, call_type: str | None - ) -> list[object]: + ) -> tuple[Message, ...]: """ Get the messages from the response object @@ -344,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): @@ -357,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, @@ -375,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 @@ -485,17 +707,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]: - 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, object]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload @@ -524,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 @@ -647,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: Sequence[object]) -> list[dict[str, object]]: - """ - 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, object]: - """ - 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, object]] = {} - 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, object]: - """ - Extract tool call information from both input messages and response for Datadog metadata. - """ - tool_call_metadata: Final[dict[str, object]] = {} - - 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/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/prometheus.py b/litellm/integrations/prometheus.py index add91033ff3..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,6 +59,7 @@ 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 @@ -476,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 ######################################## @@ -1475,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, @@ -2002,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, 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/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 207e024ce0b..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 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 eb223c2f988..9a6fb11f978 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2959,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: @@ -3005,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") diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2fca19dc8d1..ff46440ff5c 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -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 diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index e6402e8c1bd..ba59e3fa997 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4957,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/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/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/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/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 7fefeaeaf04..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, @@ -1149,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) @@ -1613,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/common_utils.py b/litellm/llms/bedrock/common_utils.py index 30a77d57f24..66ee5f10679 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -816,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 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/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index f3e41c054ef..9afc6331d96 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -20,9 +20,9 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, - flatten_top_level_schema_combinators, 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, @@ -70,19 +70,6 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) -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 - - class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Reference: https://platform.openai.com/docs/api-reference/chat/create @@ -466,7 +453,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ): 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 + tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools ] return MappingProxyType({"tools": flattened}) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index ed628f55350..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,7 +56,10 @@ 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 @@ -1005,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/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/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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55618d9f772..a3cfb300ea6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -38556,12 +38556,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" }, 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/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/_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 c549f48126e..18714256a8f 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/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/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/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/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0237d82a0d9..30526d30dc5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -31,6 +31,7 @@ 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 @@ -215,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. @@ -2709,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 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/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ff306eb65f1..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( 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 70bf8fd0554..77a80ea0052 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 @@ -9866,6 +9867,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 @@ -12481,11 +12511,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: @@ -14962,17 +15002,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() @@ -14987,7 +15035,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, @@ -15006,7 +15054,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, @@ -16434,6 +16482,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/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f4d8fd7d906..91d2ece7a51 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -120,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) @@ -134,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), @@ -258,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: 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/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index de802b95086..82abac3e772 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1023,7 +1023,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, @@ -1090,7 +1090,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, @@ -1274,10 +1274,10 @@ def _add_text_like_part_events( ) -def _build_synthetic_response_events( +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() diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..23d8907fb49 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 = [] @@ -8352,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 ) @@ -8781,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 @@ -9403,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) @@ -12087,6 +12226,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 @@ -12202,6 +12342,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_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/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/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/utils.py b/litellm/types/utils.py index 3a0883b6607..5783a39b30c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3636,6 +3636,8 @@ all_litellm_params = ( "client", "rpm", "tpm", + "default_api_key_rpm_limit", + "default_api_key_tpm_limit", "itpm", "otpm", "max_parallel_requests", diff --git a/litellm/utils.py b/litellm/utils.py index 2fdc4a45148..252b6756937 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4876,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]] = [ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55618d9f772..a3cfb300ea6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -38556,12 +38556,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" }, 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/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/models.py b/tests/e2e/models.py index 967144dfb16..79d9e011f7e 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -806,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"] @@ -820,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..c0596177cc1 --- /dev/null +++ b/tests/e2e/test_junit_properties.py @@ -0,0 +1,146 @@ +"""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, +) + + +class FakeMarker: + def __init__(self, name: str, *args: object) -> None: + self.name = name + self.args = args + + +class FakeItem: + """The three attributes junit_properties reads off a pytest Item.""" + + def __init__( + self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = () + ) -> None: + self.nodeid = nodeid + self.location = location + self.user_properties: list[tuple[str, str]] = [] + self._markers = markers + + def iter_markers(self, name: str): + return (marker for marker in self._markers if marker.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) -> None: + item = FakeItem( + "logging/test_x.py::TestFoo::test_bar", + ("logging/test_x.py", 40, "TestFoo.test_bar"), + (FakeMarker("covers", "LOG-1", "LOG-2"),), + ) + assert result_properties(item) == ( + ("package", "logging"), + ("covers", "LOG-1,LOG-2"), + ("source", "tests/e2e/logging/test_x.py:41"), + ) + + def test_attach_is_idempotent(self) -> None: + """Collection can run the hook more than once; a second pass must not + double the entries in the report.""" + item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar")) + 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/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/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/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/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/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 64c96b575c5..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 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/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/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/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 4f53d3481de..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", ) 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/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/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/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/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/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_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bcda1b8b61d..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,8 +5584,9 @@ 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: ") :]) @@ -5592,6 +5595,156 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca 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 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/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_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a81c6b4c656..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 @@ -17489,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/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/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/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/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_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_router.py b/tests/test_litellm/test_router.py index 44c1cdbff06..84f6344be35 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11595,3 +11595,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 7b7a962bf00..30b265905f3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -2281,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/type-discipline-budget.json b/type-discipline-budget.json index 3d2e97d55a5..52cb9628252 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22367 + "limit": 22364 }, "LIT002": { "limit": 26777 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/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index a504d75bb63..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 @@ -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 } } }); @@ -99,6 +108,22 @@ describe("ModelsAndEndpointsPage", () => { 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", () => { 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/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/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/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d8a955b719c..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); 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 318adcce369..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 @@ -548,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/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/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/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/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/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e944062e15e..6f044fec3f3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -13289,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; @@ -22962,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 */ @@ -25435,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; /** @@ -34921,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 */ @@ -55852,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); };