From 7abaf4edb2458b46ac4014beecd0e778ab5a9c89 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 12:30:42 -0700 Subject: [PATCH 001/154] fix(cost-map): add the Vertex shutdown date to gemini-2.5-flash-native-audio (#43024) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8f38866387a..dad89fc58a8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68565,6 +68565,7 @@ "source": "https://api.together.ai/v1/models" }, "vertex_ai/gemini-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8f38866387a..dad89fc58a8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68565,6 +68565,7 @@ "source": "https://api.together.ai/v1/models" }, "vertex_ai/gemini-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai", From c9e8a04139af587d3280ea5fb248744e3785f500 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 12:35:34 -0700 Subject: [PATCH 002/154] feat(vertex): native batch JSONL passthrough with cost tracking (#42810) * feat(vertex): native batch JSONL passthrough with cost tracking Add a per-request `passthrough=true` multipart field on `POST /v1/files` (and the same kwarg on `litellm.create_file`) that uploads a native Vertex AI batch JSONL to the deployment's GCS bucket unchanged, so rows using `googleSearch` and other Gemini-only features run as written and the output, `groundingMetadata` included, comes back untouched. Passthrough is sticky through the GCS object path (`litellm-vertex-files/passthrough/...`), so batch create and output retrieval inherit it without new state. Native output rows are costed from their `usageMetadata` with the deployment's model and model_info, in the polling and retrieve paths and for the existing global `disable_vertex_batch_output_transformation` flag, which billed $0 before. The proxy requires the target to resolve to vertex_ai deployments only, refuses `passthrough` with a non-batch purpose, a non-default `target_storage`, or pre-call guardrails, and validates native rows on `request` instead of the OpenAI batch keys. * refactor(vertex): keep native batch row pricing inside the Vertex adapter Moves native Vertex batch row detection, response parsing, and per-row pricing from litellm/batches/batch_utils.py into litellm/llms/vertex_ai/batches/transformation.py, so batch_utils only aggregates the rows it gets back. Adds tests/test_litellm/files to the misc unit shard so the new test directory is claimed by a shard. * fix(files): say what a passthrough batch upload takes when a row is not native The missing-key 400 listed bare key names, so an OpenAI-shaped row under passthrough=true read "Each line must be a JSON object with keys request". The batch line shape now carries its own hint, and the passthrough one says a passthrough upload takes native Vertex batch rows with a request key * fix(batches): bill native Vertex embedding batch rows on the native cost path A native Vertex output row whose response holds an embedding was validated as a generateContent response, so the documented tokenCount-only shape counted as a failed row. Price embedding rows from their own usage (promptTokenCount, else tokenCount) with the helper the transformed embeddings path already used, and drop the prompt-details helper nothing calls anymore. * fix(batches): keep modality batch rates on native Vertex embedding rows An embedding row that carries usageMetadata was billed from promptTokenCount alone, so its promptTokensDetails no longer reached the audio, image, and video batch rates the way it did before the native cost path. Run every row with usageMetadata through the Gemini usage parser and keep the flat tokenCount fallback for embedding rows without it. * fix(batches): price native Vertex batch rows by modelVersion under a wildcard deployment A `vertex_ai/*` deployment hands the batch cost path `*` as the deployment model, which no cost map resolves, so every native (passthrough or flag-on) row was billed at $0. A wildcard deployment model now defers to the row's own `modelVersion`, the way the transformed path already prices by the row's `model`. Also moves the native passthrough tests under tests/test_litellm, the tree codecov reads, and covers the raw upload chunking, the embedding output translation, the unpriceable-row path, and the flag-on dispatch. * fix(batches): keep explicit deployment prices for native Vertex rows without a modelVersion Under a wildcard deployment a native batch row that carries no modelVersion (an embedding row, or a generateContent row Vertex returned without one) was billed at $0 even when the deployment's model_info sets explicit batch prices, because the cost calculator was never called. The row now falls back to the wildcard name, which the cost calculator prices from the explicit model_info, and only a row with neither a modelVersion nor a deployment model is billed at $0 with the warning --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + litellm/batches/batch_utils.py | 131 +++--- litellm/files/main.py | 16 + .../llms/vertex_ai/batches/transformation.py | 147 +++++-- .../llms/vertex_ai/files/transformation.py | 127 ++++-- .../batch_file_validation.py | 35 +- .../openai_files_endpoints/files_endpoints.py | 109 ++++- litellm/router.py | 1 + litellm/router_utils/batch_utils.py | 6 +- tests/e2e/batches/test_batches_e2e.py | 112 +++++ .../llm_nonconversational.yaml | 3 + tests/e2e/coverage_registry/schema.py | 1 + tests/e2e/e2e_http.py | 7 + .../test_router_batch_utils.py | 1 + .../test_litellm/batches/test_batch_utils.py | 387 ++++++++++++++++++ tests/test_litellm/files/__init__.py | 0 tests/test_litellm/files/test_main.py | 71 ++++ .../vertex_ai/batches/test_transformation.py | 38 +- .../llms/vertex_ai/files/__init__.py | 0 .../vertex_ai/files/test_transformation.py | 310 ++++++++++++++ .../test_files_batch_file_validation.py | 38 +- .../test_files_endpoint.py | 238 +++++++++++ tests/test_litellm/test_router.py | 33 ++ tests/unit/batches/test_batch_utils.py | 17 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 + 25 files changed, 1677 insertions(+), 167 deletions(-) create mode 100644 tests/test_litellm/batches/test_batch_utils.py create mode 100644 tests/test_litellm/files/__init__.py create mode 100644 tests/test_litellm/files/test_main.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/__init__.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_transformation.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4580ad17a19..686bbc89467 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -108,6 +108,7 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/endpoints + tests/test_litellm/files tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/messages diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 7209ac6a1e7..819a279a43c 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -11,7 +11,10 @@ from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_ from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.llms.bedrock.batches.transformation import titan_embedding_usage_from_batch_output -from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details +from litellm.llms.vertex_ai.batches.transformation import ( + is_native_vertex_batch_output_row, + native_vertex_batch_row_stats, +) from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -31,6 +34,20 @@ class BatchCostUsageResult: _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) + + +def _uses_native_vertex_output( + custom_llm_provider: str, + model_name: str | None, + first_row: Mapping[str, object] | None, +) -> bool: + if custom_llm_provider != "vertex_ai": + return False + if model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False): + return True + return first_row is not None and is_native_vertex_batch_output_row(first_row) + + _TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) @@ -66,12 +83,9 @@ async def calculate_batch_cost_and_usage( deployment-specific pricing (e.g. input_cost_per_token_batches) is used instead of the global cost map. """ - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) - ): - return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) + first_row: Final = file_content_dictionary[0] if file_content_dictionary else None + if _uses_native_vertex_output(custom_llm_provider, model_name, first_row): + return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, @@ -126,11 +140,11 @@ async def _handle_completed_batch( ) output_file_result: Final = ( - calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) + calculate_vertex_ai_batch_cost_and_usage( + _iter_batch_output_entries(file_content), model_name, model_info=model_info + ) + if _uses_native_vertex_output( + custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None) ) else _aggregate_batch_cost_usage_models( entries=_iter_batch_output_entries(file_content), @@ -332,69 +346,36 @@ def _aggregate_batch_cost_usage_models( def calculate_vertex_ai_batch_cost_and_usage( - vertex_ai_batch_responses: list[dict], + vertex_ai_batch_responses: Iterable[dict], model_name: str | None = None, + model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: """ - Calculate both cost and usage from raw Vertex AI batch responses. - - Used only when ``litellm.disable_vertex_batch_output_transformation = True``. - In that case the GCS predictions.jsonl is returned as-is, with each line in - the native Vertex format: - - {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} - - usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. - - A row with no ``response`` is counted as failed - the same signal already - used to skip it from cost/usage aggregation, since Vertex batch prediction - output doesn't establish a distinct error shape in this (non-default) path. + Cost and usage of a native Vertex predictions.jsonl, one + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent row or `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding row per line. `model_name` (the deployment model) prices every row, else each row's own + `modelVersion` does; a row without a usable response counts as failed. """ from litellm.cost_calculator import batch_cost_calculator + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig - total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_tokens = 0 - prompt_tokens = 0 - completion_tokens = 0 - successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - actual_model_name: Final = model_name or "gemini-2.0-flash-001" - - for response in vertex_ai_batch_responses: - response_body = response.get("response") - if response_body is None: - failed_requests += 1 - continue - successful_requests += 1 - - usage_metadata = response_body.get("usageMetadata", {}) - _prompt = usage_metadata.get("promptTokenCount", 0) or 0 - _completion = usage_metadata.get("candidatesTokenCount", 0) or 0 - _total = usage_metadata.get("totalTokenCount", 0) or (_prompt + _completion) - - line_usage = Usage( - prompt_tokens=_prompt, - completion_tokens=_completion, - total_tokens=_total, - prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata), + row_stats: Final = tuple( + native_vertex_batch_row_stats( + row, + model_name, + model_info=model_info, + calculate_usage=VertexGeminiConfig._calculate_usage, + cost_calculator=batch_cost_calculator, ) - - try: - p_cost, c_cost = batch_cost_calculator( - usage=line_usage, - model=actual_model_name, - custom_llm_provider="vertex_ai", - ) - total_prompt_cost += p_cost - total_completion_cost += c_cost - except Exception as e: - verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) - - prompt_tokens += _prompt - completion_tokens += _completion - total_tokens += _total - + for row in vertex_ai_batch_responses + ) + priced: Final = tuple(stats for stats in row_stats if stats is not None) + total_prompt_cost: Final = sum(stats.prompt_cost for stats in priced) + total_completion_cost: Final = sum(stats.completion_cost for stats in priced) + prompt_tokens: Final = sum(stats.usage.prompt_tokens for stats in priced) + completion_tokens: Final = sum(stats.usage.completion_tokens for stats in priced) + total_tokens: Final = sum(stats.total_tokens for stats in priced) total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", @@ -402,8 +383,8 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens, completion_tokens, total_tokens, - successful_requests, - failed_requests, + len(priced), + len(row_stats) - len(priced), ) return BatchCostUsageResult( @@ -413,9 +394,13 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ), - models=[actual_model_name], - successful_requests=successful_requests, - failed_requests=failed_requests, + models=( + [model_name] + if model_name + else list(dict.fromkeys(stats.model for stats in priced if stats.model is not None)) + ), + successful_requests=len(priced), + failed_requests=len(row_stats) - len(priced), prompt_cost=total_prompt_cost, completion_cost=total_completion_cost, ) diff --git a/litellm/files/main.py b/litellm/files/main.py index e0804244ff7..72832aeccc9 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -176,6 +176,22 @@ def create_file( if logging_obj is None: raise ValueError("logging_obj is required") client: Final = kwargs.get("client") + if litellm_params_dict.get("passthrough") is True and ( + custom_llm_provider != "vertex_ai" or purpose != "batch" + ): + raise litellm.exceptions.BadRequestError( + message=( + "`passthrough=True` uploads the file bytes unchanged for a native Vertex AI batch, so it needs " + f"custom_llm_provider='vertex_ai' and purpose='batch', got '{custom_llm_provider}' and '{purpose}'." + ), + model="n/a", + llm_provider=custom_llm_provider or "n/a", + response=httpx.Response( + status_code=400, + content="passthrough needs a vertex_ai batch", + request=httpx.Request(method="create_file", url="https://github.com/BerriAI/litellm"), + ), + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f5f1ab2068a..a7dbb058465 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,11 @@ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Final, Protocol from urllib.parse import unquote +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, @@ -9,35 +13,128 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper +from litellm.types.llms.vertex_ai import GenerateContentResponseBody +from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage + +_NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody) -def vertex_prompt_tokens_details( - usage_metadata: Mapping[str, object], -) -> PromptTokensDetailsWrapper | None: - raw_details: Final = usage_metadata.get("promptTokensDetails") - if not isinstance(raw_details, list): - return None +def _int_field(mapping: Mapping[str, object], key: str) -> int: + value: Final = mapping.get(key) + if isinstance(value, int): + return value + return int(value) if isinstance(value, str) and value.isdigit() else 0 - def _normalize(detail: object) -> tuple[str, int] | None: - if not isinstance(detail, Mapping): + +def vertex_embedding_prompt_token_count(vertex_response: Mapping[str, object]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata: Final = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return _int_field(usage_metadata, "promptTokenCount") + return _int_field(vertex_response, "tokenCount") + + +def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object]) -> bool: + return isinstance(response_body.get("embedding"), dict) + + +def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool: + return isinstance(row.get("request"), dict) + + +class NativeVertexBatchCostCalculator(Protocol): + def __call__( + self, + usage: Usage, + model: str, + custom_llm_provider: str | None = None, + model_info: ModelInfo | None = None, + ) -> tuple[float, float]: ... + + +@dataclass(frozen=True, slots=True) +class NativeVertexBatchRowStats: + usage: Usage + total_tokens: int + model: str | None + prompt_cost: float + completion_cost: float + + +def _native_vertex_row_usage( + response_body: Mapping[str, object], + calculate_usage: Callable[[GenerateContentResponseBody], Usage], +) -> Usage | None: + if "usageMetadata" not in response_body: + if not is_vertex_embedding_batch_output_response(response_body): return None - modality: Final = detail.get("modality") - token_count: Final = detail.get("tokenCount") - if not isinstance(modality, str) or not isinstance(token_count, int): - return None - return modality.upper(), token_count - - parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) - normalized: Final = tuple(detail for detail in parsed_details if detail is not None) - if len(normalized) != len(parsed_details): + prompt_tokens: Final = vertex_embedding_prompt_token_count(response_body) + return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens) + try: + completion_response: Final = _NATIVE_VERTEX_RESPONSE.validate_python(response_body) + except ValidationError as e: + verbose_logger.debug("vertex_ai batch row response is not a GenerateContentResponse: %s", str(e)) return None + return calculate_usage(completion_response) - return PromptTokensDetailsWrapper( - text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), - audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), - image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), - video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + +def native_vertex_batch_row_stats( + row: Mapping[str, object], + model_name: str | None, + *, + model_info: ModelInfo | None, + calculate_usage: Callable[[GenerateContentResponseBody], Usage], + cost_calculator: NativeVertexBatchCostCalculator, +) -> NativeVertexBatchRowStats | None: + """ + Usage and cost of one native Vertex predictions.jsonl row, a + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`). + `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` + does, else the wildcard name so explicit deployment prices still apply; a row without a response, a + generateContent row without `response.usageMetadata`, and a row whose response fails validation are + None (failed). + """ + response_body: Final = row.get("response") + if not isinstance(response_body, dict): + return None + usage: Final = _native_vertex_row_usage(response_body, calculate_usage) + if usage is None: + return None + total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) + model_version: Final = response_body.get("modelVersion") + deployment_model: Final = model_name if model_name and "*" not in model_name else None + model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) + if model is None: + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed: the row has no modelVersion and the batch has no deployment model" + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=None, prompt_cost=0.0, completion_cost=0.0 + ) + try: + prompt_cost, completion_cost = cost_calculator( + usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info + ) + except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed. model=%s error=%s", + model, + str(e), + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=0.0, completion_cost=0.0 + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=prompt_cost, completion_cost=completion_cost ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 80d32289c94..789b36ef3d0 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mappin from contextlib import aclosing from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, TypedDict +from typing import IO, Any, Final, TypedDict from urllib.parse import quote, unquote import httpx @@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFileUploadStream, LiteLLMLoggingObj, ) +from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, @@ -56,6 +57,7 @@ from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContent, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, @@ -87,6 +89,8 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN: Final = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") _JSONL_NEWLINE: Final = b"\n" _BATCH_OUTPUT_FIRST_ROW_PEEK_LIMIT_BYTES: Final = 32 * 1024 * 1024 +_PASSTHROUGH_MANAGED_GCS_PREFIX: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}passthrough/" +_RAW_UPLOAD_CHUNK_BYTES: Final = 1024 * 1024 class _GcsObjectMetadataJson(TypedDict, total=False): @@ -418,19 +422,6 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[st return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) -def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> int: - """ - Prompt tokens billed for one Vertex Gemini Embedding batch row. - - Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as - a fallback. - """ - usage_metadata = vertex_response.get("usageMetadata") - if isinstance(usage_metadata, Mapping): - return int(usage_metadata.get("promptTokenCount") or 0) - return int(vertex_response.get("tokenCount") or 0) - - def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...], @@ -471,7 +462,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( ) responses = tuple(row["response"] for row in vertex_output_rows) - token_count = sum(_embedding_prompt_token_count(response) for response in responses) + token_count = sum(vertex_embedding_prompt_token_count(response) for response in responses) body = EmbeddingResponse( model=model or "", data=[ @@ -528,6 +519,16 @@ def _model_from_managed_gcs_url(url: str) -> str | None: return match.group(1) if match else None +def is_passthrough_managed_gcs_url(url: str) -> bool: + decoded_url: Final = unquote(url) + managed_prefix_start: Final = decoded_url.find(VERTEX_AI_MANAGED_GCS_PREFIX) + return managed_prefix_start >= 0 and decoded_url.startswith(_PASSTHROUGH_MANAGED_GCS_PREFIX, managed_prefix_start) + + +def is_passthrough_batch_upload(create_file_data: Mapping[str, object], litellm_params: Mapping[str, object]) -> bool: + return create_file_data.get("purpose") == "batch" and litellm_params.get("passthrough") is True + + def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: """ Whether an OpenAI batch JSONL line targets the embeddings endpoint. @@ -791,6 +792,58 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): return self._iter_vertex_jsonl_chunks() +def _read_chunk_as_bytes(handle: IO[bytes]) -> bytes: + chunk: Final[bytes | str] = handle.read(_RAW_UPLOAD_CHUNK_BYTES) + return chunk.encode("utf-8") if isinstance(chunk, str) else bytes(chunk) + + +def _iter_raw_file_chunks(file_content: FileTypes) -> Iterator[bytes]: + content: Final[FileContent | str] = file_content[1] if isinstance(file_content, tuple) else file_content + if isinstance(content, (bytes, bytearray)): + yield from ( + bytes(content[offset : offset + _RAW_UPLOAD_CHUNK_BYTES]) + for offset in range(0, len(content), _RAW_UPLOAD_CHUNK_BYTES) + ) + return + if isinstance(content, str): + yield content.encode("utf-8") + return + if isinstance(content, PathLike): + with open(str(content), "rb") as handle: + yield from iter(lambda: handle.read(_RAW_UPLOAD_CHUNK_BYTES), b"") + return + if not hasattr(content, "read"): + raise ValueError("Unsupported file content type") + seek: Final = getattr(content, "seek", None) + if seek is None: + raise ValueError( + "Batch upload file handle must be seekable; got a non-seekable " + "stream. Pass bytes, a path, or a seekable handle." + ) + seek(0) + yield from iter(lambda: _read_chunk_as_bytes(content), b"") + + +class _RawFileUploadStream(BaseFileUploadStream): + def __init__(self, file_content: FileTypes) -> None: + self._file_content = file_content + + def iter_bytes(self) -> Iterator[bytes]: + return _iter_raw_file_chunks(self._file_content) + + +def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str: + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") + prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX + return f"{prefix}{safe_model_path}/{uuid.uuid4()}" + + class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Config for VertexAI Files @@ -848,23 +901,34 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if deployment_model else openai_jsonl_content[0].get("body", {}).get("model", "") ) - endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) - model_path: Final = ( - f"endpoints/{endpoint_id}" - if endpoint_id is not None - else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") - ) - safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") - object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" - return object_name + return _managed_batch_object_name(raw_model, passthrough=False) - def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: + def _get_passthrough_gcs_object_name(self, deployment_model: str | None) -> str: + if not deployment_model: + raise VertexAIError( + status_code=400, + message=( + "Native Vertex batch passthrough uploads need the deployment model to name the GCS object, " + "since native rows carry no model: pass `target_model_names` (proxy) or `model` (SDK)." + ), + ) + return _managed_batch_object_name(deployment_model.removeprefix("vertex_ai/"), passthrough=True) + + def get_object_name( + self, + file_data: FileTypes, + purpose: str, + deployment_model: str | None = None, + passthrough: bool = False, + ) -> str: """ Get the object name for the request. Reads only the first JSONL entry (streamed) for batch files, so a large upload is never materialized just to derive the GCS object name. """ + if purpose == "batch" and passthrough: + return self._get_passthrough_gcs_object_name(deployment_model) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) @@ -922,6 +986,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): file_data, purpose, deployment_model=configured_model if isinstance(configured_model, str) else None, + passthrough=is_passthrough_batch_upload(data, litellm_params), ) if object_prefix: object_name = f"{object_prefix}/{object_name}" @@ -984,6 +1049,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if file_data is None: raise ValueError("file is required") + if is_passthrough_batch_upload(create_file_data, litellm_params): + return { + "streaming_media_upload": StreamingMediaUploadConfig( + body_stream=_RawFileUploadStream(file_data), + content_type="application/json", + ) + } + _, content_type = extract_file_metadata(file_data) if FilesAPIUtils.is_batch_jsonl_request( create_file_data=create_file_data, @@ -1164,6 +1237,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # transformation, e.g. if they consume raw `predictions.jsonl` directly. if getattr(litellm, "disable_vertex_batch_output_transformation", False): return HttpxBinaryResponseContent(response=raw_response) + if is_passthrough_managed_gcs_url(str(raw_response.request.url)): + return HttpxBinaryResponseContent(response=raw_response) # Try to transform batch output if it's a JSONL file content: Final = raw_response.content @@ -1209,7 +1284,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Everything else is passed through unchanged, including a row that fails to transform mid-stream. """ - if litellm.disable_vertex_batch_output_transformation: + if litellm.disable_vertex_batch_output_transformation or is_passthrough_managed_gcs_url(request_url): return FileContentStreamingResult(stream_iterator=stream_iterator, headers=headers) first_line, buffered = await _peek_first_jsonl_line( diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index a41bd36d510..fcd3be56ae7 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -8,10 +8,25 @@ from typing_extensions import assert_never from litellm.proxy._types import ProxyException -BATCH_LINE_REQUIRED_KEYS: Final = ("custom_id", "method", "url", "body") _MB: Final = 1024 * 1024 +@dataclass(frozen=True, slots=True) +class BatchLineShape: + required_keys: tuple[str, ...] + hint: str + + +BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("custom_id", "method", "url", "body"), + hint="Each line must be a JSON object with keys custom_id, method, url, body", +) +PASSTHROUGH_BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("request",), + hint="A passthrough upload takes native Vertex batch rows, so each line must be a JSON object with a request key", +) + + @dataclass(frozen=True, slots=True) class BatchFileTooLarge: size_bytes: int @@ -42,6 +57,7 @@ class BatchFileLineNotObject: class BatchFileMissingLineKey: line_number: int key: str + line_shape: BatchLineShape = BATCH_LINE_SHAPE BatchFileValidationFailure = ( @@ -70,20 +86,20 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: return iter(file_source) -def _check_line(line_number: int, raw_line: bytes) -> BatchFileValidationFailure | None: +def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: parsed: Final = json.loads(raw_line) except (json.JSONDecodeError, UnicodeDecodeError): return BatchFileInvalidJsonLine(line_number=line_number) if not isinstance(parsed, dict): return BatchFileLineNotObject(line_number=line_number) - missing: Final = next((key for key in BATCH_LINE_REQUIRED_KEYS if key not in parsed), None) + missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) if missing is None: return None - return BatchFileMissingLineKey(line_number=line_number, key=missing) + return BatchFileMissingLineKey(line_number=line_number, key=missing, line_shape=line_shape) -def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | None: +def _scan_lines(file_source: bytes | BinaryIO, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: content_lines: Final = ( (line_number, raw_line) for line_number, raw_line in enumerate(_iter_lines(file_source), start=1) @@ -96,7 +112,7 @@ def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | N ( failure for line_number, raw_line in chain((first_line,), content_lines) - for failure in (_check_line(line_number, raw_line),) + for failure in (_check_line(line_number, raw_line, line_shape),) if failure is not None ), None, @@ -107,6 +123,7 @@ def check_batch_file_upload( filename: str | None, file_source: bytes | BinaryIO, max_batch_file_size_mb: int | None, + line_shape: BatchLineShape = BATCH_LINE_SHAPE, ) -> BatchFileValidationFailure | None: if filename is None or not filename.lower().endswith(".jsonl"): return BatchFileWrongExtension(filename=filename or "") @@ -114,7 +131,7 @@ def check_batch_file_upload( size_bytes: Final = _file_size_bytes(file_source) if size_bytes > max_batch_file_size_mb * _MB: return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) - scan_failure: Final = _scan_lines(file_source) + scan_failure: Final = _scan_lines(file_source, line_shape) if not isinstance(file_source, bytes): file_source.seek(0) return scan_failure @@ -169,11 +186,11 @@ def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> param="file", code=400, ) - case BatchFileMissingLineKey(line_number=line_number, key=key): + case BatchFileMissingLineKey(line_number=line_number, key=key, line_shape=line_shape): raise ProxyException( message=( f"Missing required parameter: '{key}' (batch input file line {line_number}). " - f"Each line must be a JSON object with keys {', '.join(BATCH_LINE_REQUIRED_KEYS)}. " + f"{line_shape.hint}. " "The file was not forwarded to the provider." ), type="invalid_request_error", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ea2cad558c2..f5e62da5962 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -57,6 +57,8 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, ) from litellm.proxy.openai_files_endpoints.batch_file_validation import ( + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, check_batch_file_upload, raise_batch_file_validation_failure, ) @@ -207,10 +209,91 @@ def get_files_provider_config( return None +def _deployment_provider(llm_router: Router, model_id: str, team_id: str | None) -> str | None: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) + return None if credentials is None else credentials.get("custom_llm_provider") + + +def _resolves_to_vertex_deployments_only(llm_router: Router | None, model_name: str, team_id: str | None) -> bool: + if llm_router is None or _deployment_provider(llm_router, model_name, team_id) != "vertex_ai": + return False + return all( + _deployment_provider(llm_router, str(deployment["model_info"]["id"]), team_id) == "vertex_ai" + for deployment in llm_router.get_model_list(model_name=model_name, team_id=team_id) or () + if "id" in deployment.get("model_info", {}) + ) + + +def _validate_passthrough_upload( + *, + purpose: str, + target_model_names: Sequence[str], + model: str | None, + target_storage: str | None, + llm_router: Router | None, + team_id: str | None, +) -> None: + if purpose != "batch": + raise ProxyException( + message=( + "`passthrough` uploads the file bytes unchanged for a native Vertex batch, " + f"so purpose must be 'batch', got '{purpose}'." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) + if target_storage and target_storage != "default": + raise ProxyException( + message=( + "`passthrough` writes the native batch file to the Vertex AI deployment's GCS bucket, " + f"so it cannot be combined with target_storage='{target_storage}'." + ), + type="invalid_request_error", + param="target_storage", + code=400, + ) + named_deployments: Final = ( + *(("target_model_names", name) for name in target_model_names), + *((("model", model),) if model else ()), + ) + if not named_deployments: + raise ProxyException( + message=( + "`passthrough` needs the Vertex AI deployment that will run the batch, " + "since native rows carry no model: pass `target_model_names` or `model`." + ), + type="invalid_request_error", + param="target_model_names", + code=400, + ) + offending: Final = next( + ( + (param, name) + for param, name in named_deployments + if not _resolves_to_vertex_deployments_only(llm_router, name, team_id) + ), + None, + ) + if offending is None: + return + param, name = offending + raise ProxyException( + message=( + f"`passthrough` is only supported for Vertex AI deployments; '{name}' does not resolve " + "to vertex_ai deployments only." + ), + type="invalid_request_error", + param=param, + code=400, + ) + + async def _scan_batch_upload( *, file_source: bytes | BinaryIO, purpose: str, + passthrough: bool, request_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -222,6 +305,17 @@ async def _scan_batch_upload( or not proxy_logging_obj.has_pre_call_guardrails(request_metadata) ): return None + if passthrough: + raise ProxyException( + message=( + "Batch guardrails cannot scan native Vertex batch rows, so a `passthrough` upload is refused " + "when the key, team, or request has pre-call guardrails configured. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) outcome: Final = await scan_batch_input_file( file_source=file_source, request_metadata=request_metadata, @@ -458,6 +552,7 @@ async def create_file( custom_llm_provider: str = Form(default="openai"), file: UploadFile = File(...), litellm_metadata: str | None = Form(default=None), + passthrough: bool = Form(default=False), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -560,17 +655,28 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) + if passthrough: + _validate_passthrough_upload( + purpose=purpose, + target_model_names=target_model_names_list, + model=model_param, + target_storage=target_storage, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if purpose == "batch": batch_file_failure: Final = await asyncio.to_thread( check_batch_file_upload, file.filename, file_source, _MAX_BATCH_FILE_SIZE_MB_ADAPTER.validate_python(general_settings.get("max_batch_file_size_mb")), + PASSTHROUGH_BATCH_LINE_SHAPE if passthrough else BATCH_LINE_SHAPE, ) if batch_file_failure is not None: raise_batch_file_validation_failure(batch_file_failure) - data = {} + data = {"passthrough": True} if passthrough else {} # Parse expires_after if provided expires_after: FileExpiresAfter | None = None @@ -673,6 +779,7 @@ async def create_file( scan_result: Final = await _scan_batch_upload( file_source=file_source, purpose=purpose, + passthrough=passthrough, request_metadata=request_metadata, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/router.py b/litellm/router.py index d328fbbb12f..8960cd92cd8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5980,6 +5980,7 @@ class Router: replace_model_in_jsonl_bool: Final = should_replace_model_in_jsonl( purpose=purpose, + passthrough=kwargs.get("passthrough") is True, ) if replace_model_in_jsonl_bool: file = replace_model_in_jsonl( diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index be20c358202..386a2135239 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -62,15 +62,15 @@ def parse_jsonl_with_embedded_newlines(content: str) -> list[dict]: def should_replace_model_in_jsonl( purpose: OpenAIFilesPurpose, + passthrough: bool = False, ) -> bool: """ Check if the model name should be replaced in the JSONL file for the deployment model name. Azure raises an error on create batch if the model name for deployment is not in the .jsonl. + A passthrough upload keeps the caller's bytes untouched, so its rows are never rewritten. """ - if purpose == "batch": - return True - return False + return purpose == "batch" and not passthrough def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> FileTypes: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 9bb6d05bec8..6e9cf45e787 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -61,6 +61,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + proxy_error, require_successful_call, unwrap, ) @@ -1752,3 +1753,114 @@ class TestBatchTerminalState: assert (cost_row.total_tokens or 0) > 0, ( f"batch cost row has no token usage: {cost_row.total_tokens!r}" ) + + +NATIVE_VERTEX_BATCH_ROWS: Final = b"".join( + json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": text}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } + ).encode() + + b"\n" + for text in ("What is the tallest building in the world?", "Who won the last FIFA World Cup?") +) +VERTEX_BATCH_PROVIDER: Final = next(p for p in PROVIDERS if p.name == "vertex_ai") + + +class TestVertexNativePassthrough: + """`passthrough=true` on POST /v1/files uploads native Vertex batch JSONL byte for + byte (no OpenAI-to-Vertex translation, so `googleSearch` tools and the grounding + metadata they produce survive), and a batch created from that file is accepted. + + Terminal-state assertions (native output rows with groundingMetadata, the spend + row) are deliberately not here: retrieving a non-terminal batch books a $0 spend + row that blocks the real-cost row, the same reason TestBatchTerminalState polls + the list endpoint only. Those are proven by the PR's live curl proof instead. + """ + + @pytest.mark.covers( + "llm.files.vertex.native_passthrough.nonstream.works", + "llm.batches.vertex.native_passthrough.nonstream.works", + exercised_on=["files", "batches"], + ) + def test_native_jsonl_round_trips_untouched_and_starts_a_batch( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=NATIVE_VERTEX_BATCH_ROWS, + form=FileUploadForm( + purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True + ), + key=key, + ) + ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) + assert_file_object(file, provider="vertex_ai") + assert is_managed_id(file.id), f"passthrough upload must return a managed file id, got {file.id!r}" + assert file.bytes == len(NATIVE_VERTEX_BATCH_ROWS), ( + f"passthrough upload must report the caller's byte count, got {file.bytes}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", headers=client.proxy.transport.bearer(key) + ) + assert downloaded.status_code == 200, ( + f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + assert downloaded.body.encode() == NATIVE_VERTEX_BATCH_ROWS, ( + "passthrough file content must be the uploaded native rows byte for byte" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True)) + assert is_managed_id(batch.id), f"passthrough batch must be LiteLLM-managed, got {batch.id!r}" + assert batch.status in CREATED_BATCH_STATUSES, f"passthrough batch has non-transitional status {batch.status!r}" + assert batch.input_file_id == file.id + + @pytest.mark.covers("llm.files.vertex.native_passthrough_validation.nonstream.works", exercised_on=["files"]) + @pytest.mark.parametrize( + "content, form, expected_param", + [ + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", passthrough=True), + "target_model_names", + id="no-target-model", + ), + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", target_model_names=OPENAI_BATCH_MODEL, passthrough=True), + "target_model_names", + id="non-vertex-target-model", + ), + pytest.param( + render_jsonl(VERTEX_BATCH_PROVIDER.raw_model), + FileUploadForm(purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True), + "request", + id="openai-shaped-rows", + ), + ], + ) + def test_passthrough_upload_is_rejected_outside_a_native_vertex_batch( + self, + content: bytes, + form: FileUploadForm, + expected_param: str, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + result = client.upload_file(content=content, form=form, key=key) + assert isinstance(result, UnknownApiError), f"expected a 400, got {result!r}" + assert result.status_code == 400, f"expected 400, got {result.status_code}: {result.body[:300]}" + error = proxy_error(result.body) + assert error.param == expected_param, f"unexpected error param in {error!r}" + assert "passthrough" in error.message diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 3e389acc2a9..7d334ed41ff 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -21,6 +21,7 @@ - {id: llm.batches.openai_provider_fallback.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Provider-fallback raw-id scenario"} - {id: llm.batches.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Azure batches all scenarios"} - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"} +- {id: llm.batches.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "A batch created from a passthrough-uploaded native Vertex JSONL file is accepted and starts on the deployment named at upload"} - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"} - {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"} - {id: llm.batches.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create in the us-gov-west-1 partition"} @@ -46,6 +47,8 @@ - {id: llm.files.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "POST/DELETE /openai_passthrough/v1/files relay OpenAI's own file object; the dedicated prefix must not bind as a provider name on the /{provider}/v1/files route (GitHub issue #36086)"} - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} +- {id: llm.files.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "POST /v1/files with passthrough=true ships native Vertex batch JSONL (googleSearch tools and all) to GCS untouched and GET /v1/files/{id}/content returns the same bytes"} +- {id: llm.files.vertex.native_passthrough_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "passthrough=true without a Vertex target_model_names, or with OpenAI-shaped rows, is a 400 naming the offending field and nothing is uploaded"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} - {id: llm.files.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock file upload to an S3 bucket in the us-gov-west-1 partition"} - {id: llm.files.bedrock.split_s3_credentials.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: split_s3_credentials, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-8297", rationale: "Bedrock file upload, content and delete sign S3 with s3_access_key_id / s3_secret_access_key when they differ from the aws_* identity"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index e009b02b69c..e5626144fad 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -73,6 +73,7 @@ LlmCapability = Literal[ "long_context_1m", "mid_conversation_system", "multi_turn", + "native_passthrough", "pdf_input", "prompt_cache_1h", "prompt_cache_5m", diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 022caddd42a..e5d50d05c87 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -69,6 +69,7 @@ class FileUploadForm(BaseModel): purpose: str = "batch" target_model_names: str | None = None custom_llm_provider: str | None = None + passthrough: bool | None = None # ---------- Result types ---------- @@ -376,12 +377,18 @@ class ProxyErrorDetail(BaseModel): message: str type: str code: str + param: str | None = None class _ProxyErrorBody(BaseModel): error: ProxyErrorDetail +def proxy_error(body: str) -> ProxyErrorDetail: + """The proxy's own error envelope (`{"error": {message, type, param, code}}`) parsed off a rejected call.""" + return _ProxyErrorBody.model_validate_json(body).error + + def relayed_provider_rate_limit(outcome: RateLimitedError) -> ProxyErrorDetail | None: """The provider's own 429 as the proxy relayed it, or None when the 429 is the proxy's own.""" if PROVIDER_RATE_LIMIT_MARKER not in outcome.body: diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index e274ac61a01..4336185a07f 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -141,6 +141,7 @@ def test_should_replace_model_in_jsonl(): from litellm.router_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True + assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False assert should_replace_model_in_jsonl(purpose="test") is False assert should_replace_model_in_jsonl(purpose="user_data") is False diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py new file mode 100644 index 00000000000..0b2bfe9d266 --- /dev/null +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -0,0 +1,387 @@ +import json + +import pytest + +import litellm +import litellm.batches.batch_utils as bu +from litellm.types.llms.openai import Batch + +GROUNDED_USAGE_METADATA = { + "promptTokenCount": 19, + "candidatesTokenCount": 59, + "thoughtsTokenCount": 406, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 557, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], + "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], + "trafficType": "ON_DEMAND", +} +PASSTHROUGH_OUTPUT_URI = ( + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" + "predictions.jsonl" +) +UNGROUNDED_USAGE_METADATA = { + "promptTokenCount": 20, + "candidatesTokenCount": 48, + "thoughtsTokenCount": 195, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 336, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], + "trafficType": "ON_DEMAND", +} + + +def _batch(output_file_id: str) -> Batch: + return Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id=output_file_id, + ) + + +def _vertex_jsonl(rows: list[dict]) -> bytes: + return "\n".join(json.dumps(row) for row in rows).encode() + + +def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": custom_id, + "body": { + "id": f"chatcmpl-{custom_id}", + "object": "chat.completion", + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + }, + }, + "error": None, + } + + +def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): + candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} + response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} + return { + "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, + "status": "", + "response": {**response, **({"modelVersion": model_version} if model_version else {})}, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + + +def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: + import litellm.cost_calculator as cc + + calls: list = [] + + def _calc(**kw): + calls.append(kw) + return (prompt_cost, completion_cost) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + return calls + + +def test_vertex_native_cost_bills_embedding_rows(monkeypatch): + monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) + rows = [ + { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, + }, + { + "key": "id_2", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, + }, + {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") + + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) + assert result.cost == pytest.approx(5 * 1e-7) + assert result.models == ["gemini-embedding-2"] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), + ] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.models == ["gemini-2.5-flash"] + assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} + + +@pytest.mark.asyncio +async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.successful_requests == 1 + + +@pytest.mark.asyncio +async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="openai", + ) + + assert result.successful_requests == 0 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(raw_rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + result = await bu._handle_completed_batch( + _batch(PASSTHROUGH_OUTPUT_URI), + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert result.cost == pytest.approx(1.0) + assert result.usage.total_tokens == 557 + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) + ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) + + result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") + + grounded_usage, ungrounded_usage = (call["usage"] for call in calls) + assert grounded_usage.prompt_tokens == 19 + assert grounded_usage.completion_tokens == 59 + 406 + assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 + assert ungrounded_usage.prompt_tokens == 20 + 73 + assert ungrounded_usage.completion_tokens == 48 + 195 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 19 + 93, + 465 + 243, + 557 + 336, + ) + + +def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.cost == pytest.approx(1.5) + assert result.successful_requests == 3 + assert result.usage.total_tokens == 557 + 336 + 336 + + +def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): + _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, + {"request": {"contents": []}, "response": {"candidates": []}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 2) + assert result.usage.total_tokens == 557 + + +def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert result.models == ["gemini-2.5-flash"] + assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) + assert calls == [] + + +def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + bu.calculate_vertex_ai_batch_cost_and_usage( + [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + "gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6} + + await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert [call["model"] for call in calls] == ["gemini-2.5-flash"] + assert result.models == ["gemini-2.5-flash"] + + +def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 1) + assert result.usage.total_tokens == 557 + assert len(calls) == 1 + + +@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) +def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + + +def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") + without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) + + twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) + both = bu.calculate_vertex_ai_batch_cost_and_usage( + [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info + ) + + assert twin.cost > 0 + assert both.cost == pytest.approx(2 * twin.cost) + assert (both.successful_requests, both.failed_requests) == (2, 0) + + +def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): + import litellm.cost_calculator as cc + + def _calc(**kw): + if kw["model"] == "gemini-unpriced": + raise ValueError("no pricing") + return (0.5, 0.25) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert result.cost == pytest.approx(0.75) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] + + +@pytest.mark.asyncio +async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert calls == [] + assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/test_litellm/files/__init__.py b/tests/test_litellm/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/files/test_main.py b/tests/test_litellm/files/test_main.py new file mode 100644 index 00000000000..2704bdfb6ff --- /dev/null +++ b/tests/test_litellm/files/test_main.py @@ -0,0 +1,71 @@ +from typing import Final +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +NATIVE_VERTEX_ROWS: Final = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "Who won the 2024 Tour de France?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building in Tokyo?"}]}],' + b' "tools": [{"googleSearch": {}}]}}\n' +) + + +@pytest.mark.parametrize( + "custom_llm_provider, purpose", + [("openai", "batch"), ("vertex_ai", "assistants")], + ids=["non-vertex-provider", "non-batch-purpose"], +) +def test_create_file_passthrough_is_rejected_outside_a_vertex_batch(custom_llm_provider, purpose): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.create_file( + file=("batch.jsonl", b'{"request": {"contents": []}}\n', "application/jsonl"), + purpose=purpose, + custom_llm_provider=custom_llm_provider, + passthrough=True, + api_key="sk-test", + api_base="http://127.0.0.1:9", + ) + + assert "vertex_ai" in str(exc_info.value) + assert "batch" in str(exc_info.value) + + +def _gcs_upload_transport(uploads: list[httpx.Request]) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + uploads.append(request) + object_name: Final = parse_qs(urlparse(str(request.url)).query)["name"][0] + return httpx.Response( + 200, + json={ + "id": f"my-bucket/{object_name}/1758585600000000", + "name": object_name, + "size": str(len(request.read())), + "timeCreated": "2026-09-23T00:00:00.000Z", + }, + ) + + return httpx.MockTransport(respond) + + +def test_create_file_passthrough_kwarg_ships_native_rows_byte_for_byte_under_the_passthrough_prefix(): + uploads: Final[list[httpx.Request]] = [] + file_object = litellm.create_file( + file=("batch.jsonl", NATIVE_VERTEX_ROWS, "application/jsonl"), + purpose="batch", + custom_llm_provider="vertex_ai", + passthrough=True, + model="vertex_ai/gemini-2.5-flash", + gcs_bucket_name="my-bucket", + api_key="test-token", + client=HTTPHandler(client=httpx.Client(transport=_gcs_upload_transport(uploads))), + ) + (upload,) = uploads + object_name: Final = parse_qs(urlparse(str(upload.url)).query)["name"][0] + assert upload.read() == NATIVE_VERTEX_ROWS + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + assert file_object.id == f"gs://my-bucket/{object_name}" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index e6126b02790..ae045edbec5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -19,7 +19,6 @@ import pytest from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, - vertex_prompt_tokens_details, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 VertexAIError, @@ -41,27 +40,6 @@ ENDPOINT_INPUT_FILE = ( ) -def test_vertex_prompt_tokens_details_rejects_malformed_details(): - assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None - assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None - assert ( - vertex_prompt_tokens_details( - { - "promptTokensDetails": [ - {"modality": "AUDIO", "tokenCount": 1}, - "malformed", - ] - } - ) - is None - ) - - -# =========================================================================== # -# transform_openai_batch_request_to_vertex_ai_batch_request -# =========================================================================== # - - def test_transform_openai_request_builds_full_vertex_job(): with patch( "litellm.llms.vertex_ai.batches.transformation.uuid.uuid4", @@ -477,3 +455,19 @@ def test_list_response_none_jobs_treated_as_empty(): out = T.transform_vertex_ai_batch_list_response_to_openai_list_response({"batchPredictionJobs": None}) assert out["data"] == [] assert out["first_id"] is None + + +PASSTHROUGH_INPUT_FILE = ( + "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1" +) + + +def test_get_model_from_passthrough_gcs_file(): + assert T._get_model_from_gcs_file(PASSTHROUGH_INPUT_FILE) == "publishers/google/models/gemini-2.5-flash" + + +def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_input(): + assert ( + T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE) + == "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/__init__.py b/tests/test_litellm/llms/vertex_ai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/files/test_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py new file mode 100644 index 00000000000..6958576b0c8 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py @@ -0,0 +1,310 @@ +import io +import json +import urllib.parse +from pathlib import Path +from unittest.mock import MagicMock +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig, is_passthrough_managed_gcs_url + +NATIVE_VERTEX_ROW = json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": "What is the tallest building in the world?"}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } +).encode() +NATIVE_VERTEX_JSONL = NATIVE_VERTEX_ROW + b"\n" + NATIVE_VERTEX_ROW + b"\n" +OPENAI_BATCH_JSONL = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}]}}\n' +) +PASSTHROUGH_OBJECT = ( + "litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +) +TRANSFORMED_OBJECT = "litellm-vertex-files/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +@pytest.fixture +def config() -> VertexAIFilesConfig: + return VertexAIFilesConfig() + + +def _gcs_media_url(object_name: str) -> str: + return ( + f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{urllib.parse.quote(object_name, safe='')}?alt=media" + ) + + +def _native_output_jsonl() -> bytes: + return ( + json.dumps( + { + "request": json.loads(NATIVE_VERTEX_ROW)["request"], + "status": "", + "response": { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "The Burj Khalifa."}]}, + "finishReason": "STOP", + "groundingMetadata": {"webSearchQueries": ["tallest building in the world"]}, + } + ], + "modelVersion": "gemini-2.5-flash", + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 48, "totalTokenCount": 68}, + }, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + ).encode() + + b"\n" + ) + + +def _upload_chunks(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> list[bytes]: + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params=litellm_params, + ) + return list(body["streaming_media_upload"]["body_stream"].iter_bytes()) + + +def _upload_body_bytes(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> bytes: + return b"".join(_upload_chunks(config, file, litellm_params)) + + +class TestPassthroughBatchUpload: + """`passthrough=True` on a batch upload ships the caller's native Vertex JSONL + to GCS byte for byte, filed under a `passthrough/` object path so the batch + output that lands beside it is recognized and returned untouched as well.""" + + def _upload_url(self, config, litellm_params, file, purpose="batch") -> str: + return config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params=litellm_params, + data={"file": file, "purpose": purpose}, + ) + + def test_passthrough_object_is_filed_under_passthrough_prefix_named_by_deployment_model(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + + def test_passthrough_upload_without_deployment_model_is_rejected(self, config): + with pytest.raises(VertexAIError) as exc_info: + self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + assert exc_info.value.status_code == 400 + assert "target_model_names" in exc_info.value.message + + def test_passthrough_flag_does_not_ship_a_non_batch_upload_raw(self, config): + result = config.transform_create_file_request( + model="", + create_file_data={"file": ("notes.txt", b"plain text", "text/plain"), "purpose": "user_data"}, + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "passthrough": True}, + ) + assert result == b"plain text" + + def test_passthrough_flag_is_ignored_for_non_batch_purposes(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("notes.txt", b"plain text", "text/plain"), + purpose="user_data", + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/uploads/") + assert "passthrough" not in object_name + + @pytest.mark.parametrize( + "file", + [ + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + NATIVE_VERTEX_JSONL, + ("batch.jsonl", io.BytesIO(NATIVE_VERTEX_JSONL), "application/jsonl"), + ("batch.jsonl", NATIVE_VERTEX_JSONL.decode(), "application/jsonl"), + ], + ids=["bytes-tuple", "bare-bytes", "handle-tuple", "text-tuple"], + ) + def test_passthrough_upload_body_is_the_callers_bytes(self, config, file): + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params={"passthrough": True}, + ) + stream = body["streaming_media_upload"]["body_stream"] + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert body["streaming_media_upload"]["content_type"] == "application/json" + + def test_passthrough_upload_streams_a_large_handle_in_bounded_chunks(self, config): + content = NATIVE_VERTEX_ROW * (3 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + chunks = _upload_chunks( + config, ("batch.jsonl", io.BytesIO(content), "application/jsonl"), {"passthrough": True} + ) + assert len(chunks) >= 3 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_streams_a_path_in_bounded_chunks(self, config, tmp_path: Path): + content = NATIVE_VERTEX_ROW * (2 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + batch_path = tmp_path / "batch.jsonl" + batch_path.write_bytes(content) + chunks = _upload_chunks(config, ("batch.jsonl", batch_path, "application/jsonl"), {"passthrough": True}) + assert len(chunks) >= 2 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_rejects_a_non_seekable_handle(self, config): + class _Pipe: + def read(self, size=-1): + return b"" + + with pytest.raises(ValueError, match="seekable"): + _upload_body_bytes(config, ("batch.jsonl", _Pipe(), "application/jsonl"), {"passthrough": True}) + + def test_passthrough_upload_rejects_content_that_is_neither_bytes_path_nor_handle(self, config): + with pytest.raises(ValueError, match="Unsupported file content type"): + _upload_body_bytes(config, ("batch.jsonl", 42, "application/jsonl"), {"passthrough": True}) + + def test_openai_rows_are_translated_unless_passthrough_is_set(self, config): + file = ("batch.jsonl", OPENAI_BATCH_JSONL, "application/jsonl") + translated = _upload_body_bytes(config, file, {}) + untouched = _upload_body_bytes(config, file, {"passthrough": True}) + assert untouched == OPENAI_BATCH_JSONL + assert translated != OPENAI_BATCH_JSONL + assert b'"contents"' in translated + + def test_passthrough_output_content_is_returned_untouched(self, config): + raw_jsonl = _native_output_jsonl() + + def _download(object_name: str) -> bytes: + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(object_name)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return result.response.content + + assert _download(PASSTHROUGH_OBJECT) == raw_jsonl + assert _download(f"team-a/{PASSTHROUGH_OBJECT}") == raw_jsonl + transformed = _download(TRANSFORMED_OBJECT) + assert transformed != raw_jsonl + assert json.loads(transformed.splitlines()[0])["response"]["body"]["choices"] + nested = _download(f"litellm-vertex-files/{PASSTHROUGH_OBJECT}") + assert nested != raw_jsonl + assert json.loads(nested.splitlines()[0])["response"]["body"]["choices"] + + def test_output_of_an_upload_whose_model_smuggles_the_passthrough_segment_is_still_transformed(self, config): + smuggled_model = b"litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + upload_url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket"}, + ("batch.jsonl", OPENAI_BATCH_JSONL.replace(b"gemini-2.5-flash", smuggled_model), "application/jsonl"), + ) + object_name = parse_qs(urlparse(upload_url).query)["name"][0] + raw_jsonl = _native_output_jsonl() + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(f"{object_name}/predictions.jsonl")), + ) + + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + + assert object_name.startswith("litellm-vertex-files/litellm-vertex-files/passthrough/") + assert json.loads(result.response.content.splitlines()[0])["response"]["body"]["choices"] + + @pytest.mark.parametrize( + "url, expected", + [ + (f"gs://my-bucket/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/team-a/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/litellm-vertex-files/{PASSTHROUGH_OBJECT}", False), + (_gcs_media_url(f"team-a/sub/{PASSTHROUGH_OBJECT}"), True), + (_gcs_media_url(f"litellm-vertex-files/publishers/google/models/x/{PASSTHROUGH_OBJECT}"), False), + (_gcs_media_url(TRANSFORMED_OBJECT), False), + ], + ids=["gs", "gs-prefixed", "gs-smuggled", "https-prefixed", "https-model-path-smuggled", "https-transformed"], + ) + def test_passthrough_detection_anchors_on_the_first_managed_segment(self, url, expected): + assert is_passthrough_managed_gcs_url(url) is expected + + @pytest.mark.asyncio + async def test_passthrough_output_stream_is_returned_untouched(self, config): + stream_iterator = object() + headers = {"content-type": "application/octet-stream"} + result = await config.transform_file_content_stream( + stream_iterator=stream_iterator, + headers=headers, + request_url=_gcs_media_url(f"team-a/{PASSTHROUGH_OBJECT}"), + logging_obj=MagicMock(), + litellm_params={}, + ) + assert result.stream_iterator is stream_iterator + assert result.headers == headers + + +class TestEmbeddingOutputTranslation: + EMBEDDING_OBJECT = ( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/prediction-model-1/predictions.jsonl" + ) + + def _transform(self, config: VertexAIFilesConfig, rows: list[dict]) -> list[dict]: + raw_response = httpx.Response( + status_code=200, + content="\n".join(json.dumps(row) for row in rows).encode(), + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(self.EMBEDDING_OBJECT)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return [json.loads(line) for line in result.response.content.decode().splitlines()] + + def test_embedding_rows_become_openai_batch_rows_billed_by_their_prompt_tokens(self, config): + live_row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}, + } + documented_row = { + "key": "request-2", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.5]}, "tokenCount": "3"}, + } + + live, documented = self._transform(config, [live_row, documented_row]) + + assert (live["custom_id"], live["error"], live["response"]["status_code"]) == ("request-1", None, 200) + assert live["response"]["body"]["model"] == "gemini-embedding-2" + assert live["response"]["body"]["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] + live_usage, documented_usage = (row["response"]["body"]["usage"] for row in (live, documented)) + assert (live_usage["prompt_tokens"], live_usage["total_tokens"]) == (2, 2) + assert (documented_usage["prompt_tokens"], documented_usage["total_tokens"]) == (3, 3) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py index f5542fc0446..3a73c39e177 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -4,7 +4,8 @@ import pytest from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.batch_file_validation import ( - BATCH_LINE_REQUIRED_KEYS, + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, BatchFileEmpty, BatchFileInvalidJsonLine, BatchFileLineNotObject, @@ -96,7 +97,7 @@ def test_non_object_line_rejected(): assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileLineNotObject(line_number=2) -@pytest.mark.parametrize("missing_key", BATCH_LINE_REQUIRED_KEYS) +@pytest.mark.parametrize("missing_key", BATCH_LINE_SHAPE.required_keys) def test_missing_required_key_rejected(missing_key): import json @@ -174,3 +175,36 @@ def test_failures_map_to_openai_shaped_proxy_exceptions(failure, expected_code, assert exc_info.value.param == expected_param for fragment in expected_fragments: assert fragment in exc_info.value.message + + +NATIVE_VERTEX_LINE = b'{"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}}' + + +def test_passthrough_keys_accept_native_vertex_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + NATIVE_VERTEX_LINE + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE) is None + + +def test_passthrough_keys_reject_openai_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + VALID_LINE + b"\n" + assert check_batch_file_upload( + "batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE + ) == BatchFileMissingLineKey(line_number=2, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + + +def test_default_keys_still_reject_native_vertex_rows(): + assert check_batch_file_upload("batch.jsonl", NATIVE_VERTEX_LINE, None) == BatchFileMissingLineKey( + line_number=1, key="custom_id" + ) + + +def test_passthrough_missing_key_message_says_what_a_passthrough_upload_takes(): + with pytest.raises(ProxyException) as exc_info: + raise_batch_file_validation_failure( + BatchFileMissingLineKey(line_number=3, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + ) + assert exc_info.value.param == "request" + assert "line 3" in exc_info.value.message + assert "passthrough upload takes native Vertex batch rows" in exc_info.value.message + assert "with a request key." in exc_info.value.message + assert "custom_id" not in exc_info.value.message diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 48699b47e7f..84cf4ea7c32 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -5669,3 +5669,241 @@ def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFi assert response.status_code == 200, response.text assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["custom_llm_provider"] == "mistral" + + +NATIVE_VERTEX_BATCH_LINE = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' +) + + +def _passthrough_router() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + ] + ) + + +def _setup_passthrough_upload_endpoint(monkeypatch, llm_router: Router) -> list: + """Like _setup_batch_upload_endpoint, but reads the forwarded file bytes while the spool is open.""" + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + async def fake_route_create_file(**kwargs): + upload_source = kwargs["_create_file_request"]["file"][1] + upload_source.seek(0) + forwarded_calls.append({**kwargs, "file_bytes": upload_source.read()}) + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + return forwarded_calls + + +def _upload(content: bytes, form: dict): + return client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data=form, + headers={"Authorization": "Bearer test-key"}, + ) + + +def test_create_file_passthrough_forwards_native_vertex_rows_untouched(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + content = NATIVE_VERTEX_BATCH_LINE * 2 + + try: + response = _upload(content, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["_create_file_request"]["passthrough"] is True + assert call["file_bytes"] == content + assert call["target_model_names_list"] == ["vertex-batch"] + + +def test_create_file_passthrough_rejects_rows_without_a_request(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE + VALID_BATCH_LINE, + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "request" + assert "line 2" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_without_passthrough_still_rejects_native_vertex_rows(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "custom_id" + assert forwarded_calls == [] + + +@pytest.mark.parametrize( + "form, expected_param, expected_fragment", + [ + ({"purpose": "batch", "passthrough": "true"}, "target_model_names", "target_model_names"), + ( + {"purpose": "batch", "target_model_names": "gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ( + {"purpose": "batch", "target_model_names": "vertex-batch,gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ({"purpose": "user_data", "target_model_names": "vertex-batch", "passthrough": "true"}, "passthrough", "batch"), + ( + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true", "target_storage": "s3"}, + "target_storage", + "'s3'", + ), + ( + {"purpose": "batch", "model": "gpt-3.5-turbo", "passthrough": "true"}, + "model", + "'gpt-3.5-turbo'", + ), + ], + ids=[ + "no-model", + "non-vertex-model", + "mixed-models", + "non-batch-purpose", + "target-storage", + "non-vertex-model-param", + ], +) +def test_create_file_passthrough_rejected_outside_a_vertex_batch(monkeypatch, form, expected_param, expected_fragment): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, form) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == expected_param + assert expected_fragment in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_accepts_the_model_param_as_the_deployment(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "model": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["model"] == "vertex-batch" + assert call["_create_file_request"]["passthrough"] is True + + +def test_create_file_passthrough_rejects_a_model_group_with_a_non_vertex_deployment(monkeypatch): + mixed_router = Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "vertex-batch", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "openai_api_key"}, + "model_info": {"id": "vertex-batch-openai-id"}, + }, + ] + ) + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, mixed_router) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "target_model_names" + assert "'vertex-batch'" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_batch(monkeypatch): + """Batch guardrails read OpenAI-shaped rows, so a passthrough upload on a guardrailed + key is refused rather than forwarded unscanned.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + ProxyLogging._callback_capabilities_cache.clear() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "passthrough" + assert "guardrails" in error["message"] + assert forwarded_calls == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 411377f29cf..f985b212b01 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -528,6 +528,39 @@ async def test_async_router_acreate_file_with_jsonl(): assert first_call_content == non_jsonl_content +@pytest.mark.asyncio +async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards_the_flag(): + """A passthrough batch upload must reach the provider byte for byte: the router + neither rewrites body.model to the deployment model nor drops the flag.""" + from io import BytesIO + from unittest.mock import MagicMock, patch + + jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "vertex_project": "p"}, + } + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file( + model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True + ) + forwarded = mock_acreate_file.call_args.kwargs + assert forwarded["passthrough"] is True + forwarded["file"].seek(0) + assert forwarded["file"].read() == jsonl_content + + mock_acreate_file.reset_mock() + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content)) + rewritten = mock_acreate_file.call_args.kwargs["file"] + rewritten.seek(0) + assert b'"gemini-2.5-flash"' in rewritten.read() + + @pytest.mark.asyncio async def test_async_router_acreate_file_does_not_fall_back_across_model_groups(): """A file created for batches only exists under the credentials of the model group diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index a4de8eee23c..d1572f4a7c9 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -640,7 +640,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: bu.BatchCostUsageResult( + lambda content, model, model_info=None: bu.BatchCostUsageResult( cost=9.9, usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), models=["gemini-2.0-flash-001"], @@ -671,7 +671,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: pytest.fail("raw vertex path should not run"), + lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"), ) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") @@ -735,7 +735,11 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): ) responses = [ { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}, {"fileData": {"mimeType": "audio/wav"}}]}}, "response": { + "embedding": {"values": [0.1, 0.2]}, "usageMetadata": { "promptTokenCount": 84, "candidatesTokenCount": 0, @@ -744,13 +748,14 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): {"modality": "AUDIO", "tokenCount": 64}, {"modality": "TEXT", "tokenCount": 20}, ], - } - } + }, + }, } ] result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2") + assert (result.successful_requests, result.usage.prompt_tokens) == (1, 84) assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7) @@ -1336,7 +1341,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) seen: dict = {} - def fake_vertex_calc(content, model): + def fake_vertex_calc(content, model, model_info=None): seen["content"] = content seen["model"] = model return bu.BatchCostUsageResult( @@ -1358,7 +1363,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) assert result.cost == 7.7 assert result.usage.total_tokens == 3 assert result.models == ["gemini-x"] - assert seen["content"] == raw_rows + assert list(seen["content"]) == raw_rows assert seen["model"] == "gemini-x" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2938e65cfde..052cc693f8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25526,6 +25526,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25550,6 +25555,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25574,6 +25584,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** From 1edc4ba580091d3e3d344b63d3f92d08b9b21a4d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:01:12 -0700 Subject: [PATCH 003/154] fix(logging): pass provider response headers to callbacks on every endpoint (#42824) * fix(logging): pass provider response headers to callbacks on every endpoint Custom callbacks only received kwargs["response_headers"] for chat completions. Responses, image generation and edit, speech, and transcription calls either never recorded the provider's headers or recorded them in one place and not the other. Every handler now records the provider's httpx headers on the response's hidden params as "headers" (raw) and "additional_headers" (processed, with LiteLLM's own entries winning on a clash), and the logging object derives model_call_details["response_headers"] from those hidden params before cost calculation on the non-stream and both streaming success paths, keeping a handler-set value authoritative. Binary speech responses expose their hidden params to the standard logging payload, and the sync OpenAI transcription request always fetches the raw response. * test(images): point the legacy image and speech fakes at the raw response surface Image generation now goes through the SDK's raw response so the provider headers can be read, and the speech binary response now carries hidden params. The unit fakes in the image generation, xinference, proxy provider, image edit, Vertex speech, and otel suites still pinned the old call surface and the old "no hidden params" assertion, so they read an uncalled mock or a fake response without headers. * test(images): drop the rewritten mock comments and the generated edit PNGs * test(images): move the llm-span test's image fake to the raw response surface --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 27 +++ litellm/litellm_core_utils/litellm_logging.py | 20 +- litellm/llms/custom_httpx/llm_http_handler.py | 20 +- litellm/llms/openai/openai.py | 29 ++- litellm/llms/openai/transcriptions/handler.py | 33 +-- tests/image_gen_tests/test_image_edits.py | 4 + tests/image_gen_tests/test_xinference.py | 30 ++- .../test_litellm_proxy_provider.py | 16 +- tests/llm_translation/test_openai.py | 13 +- .../otel/test_otel_v2_sources_of_truth.py | 8 +- .../litellm_core_utils/test_core_helpers.py | 68 ++++++ .../test_litellm_logging.py | 87 ++++++++ .../custom_httpx/test_llm_http_handler.py | 203 +++++++++++++++++- tests/test_litellm/llms/openai/test_openai.py | 99 ++++++++- .../test_openai_transcriptions_handler.py | 71 ++++++ .../test_non_chat_routes_open_llm_spans.py | 14 +- ...t_openai_image_generation_extra_headers.py | 50 +++-- .../text_to_speech/test_transformation.py | 1 + 18 files changed, 714 insertions(+), 79 deletions(-) create mode 100644 tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 3afa6a913b5..b095b4b12c6 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -765,3 +765,30 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo RESPONSE_COST_HEADER: cost, } hidden_params["additional_headers"] = merged + + +_HIDDEN_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_PROVIDER_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, str]) + + +def set_provider_response_headers_in_hidden_params( + response: _CarriesHiddenParams, headers: httpx.Headers | Mapping[str, str] +) -> None: + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + existing_additional_headers: Final[object] = hidden_params.get("additional_headers") + raw_headers: Final[dict[str, str]] = dict(headers) # mutable-ok: stored as the plain-dict hidden param + additional_headers: Final[dict[str, object]] = { # mutable-ok: assigned into the plain-dict hidden params + **process_response_headers(raw_headers), + **(existing_additional_headers if isinstance(existing_additional_headers, Mapping) else _NO_HEADERS), + } + hidden_params["headers"] = raw_headers + hidden_params["additional_headers"] = additional_headers + + +def get_provider_response_headers_from_hidden_params(response: object) -> Mapping[str, str] | None: + hidden_params: Final[object] = getattr(response, "_hidden_params", None) + try: + validated: Final = _HIDDEN_PARAMS_ADAPTER.validate_python(hidden_params) + return _PROVIDER_HEADERS_ADAPTER.validate_python(validated.get("headers")) + except ValidationError: + return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a6391a2ae27..3b9419d483b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -72,6 +72,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + get_provider_response_headers_from_hidden_params, is_expected_client_error, reconstruct_model_name, set_response_cost_in_hidden_params, @@ -2353,6 +2354,15 @@ class Logging(LiteLLMLoggingBaseClass): ) return logging_result + def _surface_response_headers_from_result(self, logging_result: object) -> None: + existing: Final[object] = self.model_call_details.get("response_headers") + if existing is not None: + return + headers: Final = get_provider_response_headers_from_hidden_params(logging_result) + if headers is None: + return + self.model_call_details["response_headers"] = headers + def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. @@ -2386,6 +2396,7 @@ class Logging(LiteLLMLoggingBaseClass): build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" + self._surface_response_headers_from_result(logging_result) hidden_params: Final = getattr(logging_result, "_hidden_params", {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: @@ -2788,6 +2799,7 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: verbose_logger.debug("Logging Details LiteLLM-Success Call streaming complete") self.model_call_details["complete_streaming_response"] = complete_streaming_response + self._surface_response_headers_from_result(complete_streaming_response) self.model_call_details["response_cost"] = self._response_cost_calculator( result=complete_streaming_response ) @@ -3302,6 +3314,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("Async success callbacks: Got a complete streaming response") self.model_call_details["async_complete_streaming_response"] = complete_streaming_response + self._surface_response_headers_from_result(complete_streaming_response) try: if self.model_call_details.get("cache_hit", False) is True: @@ -6362,12 +6375,15 @@ def _extract_response_obj_and_hidden_params( original_exception: Exception | None, ) -> tuple[dict, dict | None]: """Extract response_obj and hidden_params from init_response_obj.""" - hidden_params: dict | None = None + hidden_params: dict | None = ( + getattr(init_response_obj, "_hidden_params", None) + if isinstance(init_response_obj, BaseModel | HttpxBinaryResponseContent) + else None + ) if init_response_obj is None: response_obj = {} elif isinstance(init_response_obj, BaseModel): response_obj = init_response_obj.model_dump() - hidden_params = getattr(init_response_obj, "_hidden_params", None) elif isinstance(init_response_obj, dict): response_obj = init_response_obj elif isinstance(init_response_obj, HttpxBinaryResponseContent): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dba0dee38fc..4f31742aaaa 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -45,6 +45,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( @@ -1461,6 +1462,7 @@ class BaseLLMHTTPHandler: transformed: Final = provider_config.transform_audio_transcription_response( raw_response=response, ) + set_provider_response_headers_in_hidden_params(transformed, response.headers) if not provider_config.supports_subtitle_synthesis: return transformed requested_format: Final = optional_params.get("response_format") @@ -6960,11 +6962,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( + image_edit_response: Final = image_edit_provider_config.transform_image_edit_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(image_edit_response, response.headers) + return image_edit_response async def async_image_edit_handler( self, @@ -7059,11 +7063,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( + image_edit_response: Final = image_edit_provider_config.transform_image_edit_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(image_edit_response, response.headers) + return image_edit_response def image_generation_handler( self, @@ -7186,6 +7192,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), encoding=None, ) + set_provider_response_headers_in_hidden_params(model_response, response.headers) return model_response @@ -7293,6 +7300,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), encoding=None, ) + set_provider_response_headers_in_hidden_params(model_response, response.headers) return model_response @@ -12077,11 +12085,13 @@ class BaseLLMHTTPHandler: provider_config=text_to_speech_provider_config, ) - return text_to_speech_provider_config.transform_text_to_speech_response( + speech_response: Final = text_to_speech_provider_config.transform_text_to_speech_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(speech_response, response.headers) + return speech_response async def async_text_to_speech_handler( self, @@ -12176,11 +12186,13 @@ class BaseLLMHTTPHandler: provider_config=text_to_speech_provider_config, ) - return text_to_speech_provider_config.transform_text_to_speech_response( + speech_response: Final = text_to_speech_provider_config.transform_text_to_speech_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(speech_response, response.headers) + return speech_response ######################################################### ########## SKILLS API HANDLERS ########################## diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 63874ca9619..d6340d182ae 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -27,6 +27,7 @@ from litellm import LlmProviders from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES from litellm.files.types import FileContentStreamingResult +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -1404,7 +1405,6 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): organization: str | None = None, headers: dict | None = None, ): - response = None try: openai_aclient: Final = self._get_openai_client( is_async=True, @@ -1428,8 +1428,10 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) request_data: Final = {**data, "extra_headers": headers} if headers else data - response = await openai_aclient.images.generate(**request_data, timeout=timeout) - stringified_response: Final = response.model_dump() + raw_response: Final = await openai_aclient.images.with_raw_response.generate( + **request_data, timeout=timeout + ) + stringified_response: Final = raw_response.parse().model_dump() ## LOGGING logging_obj.post_call( input=prompt, @@ -1437,11 +1439,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=stringified_response, ) - return convert_to_model_response_object( + image_response: Final[ImageResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, response_type="image_generation", ) + set_provider_response_headers_in_hidden_params(image_response, raw_response.headers) + return image_response except Exception as e: ## LOGGING logging_obj.post_call( @@ -1512,9 +1516,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ## COMPLETION CALL request_data: Final = {**data, "extra_headers": headers} if headers else data - _response: Final = openai_client.images.generate(**request_data, timeout=timeout) + raw_response: Final = openai_client.images.with_raw_response.generate(**request_data, timeout=timeout) - response: Final = _response.model_dump() + response: Final = raw_response.parse().model_dump() ## LOGGING logging_obj.post_call( input=prompt, @@ -1522,11 +1526,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=response, ) - return convert_to_model_response_object( + image_response: Final[ImageResponse] = convert_to_model_response_object( response_object=response, model_response_object=model_response, response_type="image_generation", ) + set_provider_response_headers_in_hidden_params(image_response, raw_response.headers) + return image_response except OpenAIError as e: ## LOGGING logging_obj.post_call( @@ -1609,7 +1615,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): input=input, **optional_params, ) - return HttpxBinaryResponseContent(response=response.response) + speech_response: Final = HttpxBinaryResponseContent(response=response.response) + set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + return speech_response async def async_audio_speech( self, @@ -1655,8 +1663,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): input=input, **optional_params, ) - - return HttpxBinaryResponseContent(response=response.response) + speech_response: Final = HttpxBinaryResponseContent(response=response.response) + set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + return speech_response class OpenAIFilesAPI(BaseLLM): diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 701b3d30362..014251db821 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -4,11 +4,10 @@ import httpx from openai import AsyncOpenAI, OpenAI from pydantic import BaseModel -import litellm - if TYPE_CHECKING: from aiohttp import ClientSession from litellm.litellm_core_utils.audio_utils.utils import get_audio_file_name +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, @@ -31,11 +30,6 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): data: dict, timeout: float | httpx.Timeout, ): - """ - Helper to: - - call openai_aclient.audio.transcriptions.with_raw_response when litellm.return_response_headers is True - - call openai_aclient.audio.transcriptions.create by default - """ try: raw_response = await openai_aclient.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) headers: Final = dict(raw_response.headers) @@ -51,20 +45,11 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): data: dict, timeout: float | httpx.Timeout, ): - """ - Helper to: - - call openai_aclient.audio.transcriptions.with_raw_response when litellm.return_response_headers is True - - call openai_aclient.audio.transcriptions.create by default - """ try: - if litellm.return_response_headers is True: - raw_response = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response = raw_response.parse() - return headers, response - else: - response = openai_client.audio.transcriptions.create(**data, timeout=timeout) - return None, response + raw_response: Final = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) + headers: Final = dict(raw_response.headers) + response: Final = raw_response.parse() + return headers, response except Exception as e: raise e @@ -133,11 +118,12 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): "complete_input_dict": data, }, ) - _, response = self.make_sync_openai_audio_transcriptions_request( + headers, response = self.make_sync_openai_audio_transcriptions_request( openai_client=openai_client, data=data, timeout=timeout, ) + logging_obj.model_call_details["response_headers"] = headers if isinstance(response, BaseModel): stringified_response = response.model_dump() @@ -158,6 +144,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): hidden_params=hidden_params, response_type="audio_transcription", ) + set_provider_response_headers_in_hidden_params(final_response, headers) return final_response async def async_audio_transcriptions( @@ -217,12 +204,14 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): actual_model: Final = data.get("model", "whisper-1") hidden_params: Final = {"model": actual_model, "custom_llm_provider": "openai"} - return convert_to_model_response_object( + final_response: Final[TranscriptionResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription", ) + set_provider_response_headers_in_hidden_params(final_response, headers) + return final_response except Exception as e: ## LOGGING logging_obj.post_call( diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 0c2f57066e8..36fd65ba71b 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -250,6 +250,7 @@ async def test_azure_image_edit_litellm_sdk(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -370,6 +371,7 @@ async def test_openai_image_edit_cost_tracking(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -460,6 +462,7 @@ async def test_azure_image_edit_cost_tracking(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -737,6 +740,7 @@ async def test_image_edit_array_handling(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data diff --git a/tests/image_gen_tests/test_xinference.py b/tests/image_gen_tests/test_xinference.py index 3dc4fee85da..76cae593e41 100644 --- a/tests/image_gen_tests/test_xinference.py +++ b/tests/image_gen_tests/test_xinference.py @@ -24,9 +24,14 @@ async def test_xinference_image_generation(): def model_dump(self): return mock_openai_response - # Create a mock client with the images.generate method + class MockRawResponse: + headers = {} + + def parse(self): + return MockResponse() + mock_client = AsyncMock() - mock_client.images.generate = AsyncMock(return_value=MockResponse()) + mock_client.images.with_raw_response.generate = AsyncMock(return_value=MockRawResponse()) # Capture the actual arguments sent to OpenAI client captured_args = None @@ -36,9 +41,9 @@ async def test_xinference_image_generation(): nonlocal captured_args, captured_kwargs captured_args = args captured_kwargs = kwargs - return MockResponse() + return MockRawResponse() - mock_client.images.generate.side_effect = capture_generate_call + mock_client.images.with_raw_response.generate.side_effect = capture_generate_call # Mock the _get_openai_client method to return our mock client with patch.object( @@ -65,7 +70,7 @@ async def test_xinference_image_generation(): assert response.data[0].url == "https://example.com/image.png" # Validate that the OpenAI client was called with correct parameters - mock_client.images.generate.assert_called_once() + mock_client.images.with_raw_response.generate.assert_called_once() assert captured_kwargs is not None assert ( captured_kwargs["model"] == "stabilityai/stable-diffusion-3.5-large" @@ -97,9 +102,14 @@ async def test_xinference_image_generation_with_response_format(): def model_dump(self): return mock_openai_response - # Create a mock client with the images.generate method + class MockRawResponse: + headers = {} + + def parse(self): + return MockResponse() + mock_client = AsyncMock() - mock_client.images.generate = AsyncMock(return_value=MockResponse()) + mock_client.images.with_raw_response.generate = AsyncMock(return_value=MockRawResponse()) # Capture the actual arguments sent to OpenAI client captured_args = None @@ -109,9 +119,9 @@ async def test_xinference_image_generation_with_response_format(): nonlocal captured_args, captured_kwargs captured_args = args captured_kwargs = kwargs - return MockResponse() + return MockRawResponse() - mock_client.images.generate.side_effect = capture_generate_call + mock_client.images.with_raw_response.generate.side_effect = capture_generate_call # Mock the _get_openai_client method to return our mock client with patch.object( @@ -141,7 +151,7 @@ async def test_xinference_image_generation_with_response_format(): assert response.data[0].b64_json is not None # Validate that the OpenAI client was called with correct parameters - mock_client.images.generate.assert_called_once() + mock_client.images.with_raw_response.generate.assert_called_once() assert captured_kwargs is not None assert ( captured_kwargs["model"] == "stabilityai/stable-diffusion-3.5-large" diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index a10fc55ecc5..a7a2848a514 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -210,11 +210,14 @@ async def test_litellm_gateway_image_generation_direct(is_async): "created": 1, "data": [{"url": "https://example.com/image.png"}], } + mock_raw_response = MagicMock() + mock_raw_response.parse.return_value = mock_openai_response + mock_raw_response.headers = {} if is_async: # Mock the AsyncOpenAI client that gets created inside _get_openai_client mock_async_client = AsyncMock() - mock_async_client.images.generate = AsyncMock(return_value=mock_openai_response) + mock_async_client.images.with_raw_response.generate = AsyncMock(return_value=mock_raw_response) with patch( "litellm.llms.openai.openai.AsyncOpenAI", return_value=mock_async_client @@ -234,14 +237,14 @@ async def test_litellm_gateway_image_generation_direct(is_async): assert constructor_kwargs["base_url"] == "http://my-proxy" # Verify the AsyncOpenAI client was called correctly - mock_async_client.images.generate.assert_awaited_once() - call_kwargs = mock_async_client.images.generate.call_args.kwargs + mock_async_client.images.with_raw_response.generate.assert_awaited_once() + call_kwargs = mock_async_client.images.with_raw_response.generate.call_args.kwargs assert call_kwargs["model"] == "dall-e-3" assert call_kwargs["prompt"] == "A beautiful sunset over mountains" else: # Mock the sync OpenAI client that gets created inside _get_openai_client mock_sync_client = MagicMock() - mock_sync_client.images.generate.return_value = mock_openai_response + mock_sync_client.images.with_raw_response.generate.return_value = mock_raw_response with patch( "litellm.llms.openai.openai.OpenAI", return_value=mock_sync_client @@ -260,8 +263,8 @@ async def test_litellm_gateway_image_generation_direct(is_async): assert constructor_kwargs["base_url"] == "http://my-proxy" # Verify the OpenAI client was called correctly - mock_sync_client.images.generate.assert_called_once() - call_kwargs = mock_sync_client.images.generate.call_args.kwargs + mock_sync_client.images.with_raw_response.generate.assert_called_once() + call_kwargs = mock_sync_client.images.with_raw_response.generate.call_args.kwargs assert call_kwargs["model"] == "dall-e-3" assert call_kwargs["prompt"] == "A beautiful sunset over mountains" @@ -285,6 +288,7 @@ async def test_litellm_gateway_from_sdk_image_edit(is_async): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index af4ba85d58e..0488c4c68e6 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -313,7 +313,7 @@ def test_openai_max_retries_0(mock_get_openai_client): def test_openai_image_generation_forwards_organization(mock_get_openai_client): """Ensure organization flows to OpenAI client for image generation.""" - class _DummyImages: + class _DummyRawImages: def generate(self, **kwargs): # type: ignore class _Resp: def model_dump(self_inner): # minimal OpenAI ImagesResponse shape @@ -327,7 +327,16 @@ def test_openai_image_generation_forwards_organization(mock_get_openai_client): }, } - return _Resp() + class _RawResp: + headers = {} + + def parse(self_inner): + return _Resp() + + return _RawResp() + + class _DummyImages: + with_raw_response = _DummyRawImages() class _DummyClient: def __init__(self): 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 7e93d3d67a7..57f4557c6f7 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 @@ -1207,14 +1207,18 @@ def test_speech_response_without_a_byte_count_produces_no_output() -> None: def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None: import httpx + from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import _extract_response_obj_and_hidden_params from litellm.types.llms.openai import HttpxBinaryResponseContent raw: Final = httpx.Response(200, headers={"content-type": "audio/mpeg"}, content=b"\x00" * 1234) - response_obj, hidden_params = _extract_response_obj_and_hidden_params(HttpxBinaryResponseContent(raw), None) + speech: Final = HttpxBinaryResponseContent(raw) + set_provider_response_headers_in_hidden_params(speech, raw.headers) + response_obj, hidden_params = _extract_response_obj_and_hidden_params(speech, None) assert response_obj == {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 1234} - assert hidden_params is None + assert hidden_params is not None + assert hidden_params["headers"]["content-type"] == "audio/mpeg" def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far() -> None: diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 6eeea271127..2a6dd347d5f 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -2,22 +2,27 @@ import logging +import httpx import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + RESPONSE_COST_HEADER, bind_budget_reservation_to_callbacks, budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, + get_provider_response_headers_from_hidden_params, map_finish_reason, normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, + set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import ImageResponse, TranscriptionResponse class TestBudgetReservationBinding: @@ -489,3 +494,66 @@ class TestIsExpectedClientError: category=RateLimitErrorCategory.VENDOR_RATE_LIMIT, ) assert is_expected_client_error(vendor_limit) is False + + +class TestProviderResponseHeadersInHiddenParams: + def test_records_raw_headers_and_the_processed_additional_headers(self): + response = ImageResponse() + response._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: 0.04}} + + set_provider_response_headers_in_hidden_params( + response, httpx.Headers({"X-Request-Id": "req_img", "x-ratelimit-remaining-requests": "41"}) + ) + + assert response._hidden_params["headers"] == { + "x-request-id": "req_img", + "x-ratelimit-remaining-requests": "41", + } + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "req_img" + assert additional_headers["x-ratelimit-remaining-requests"] == "41" + assert additional_headers[RESPONSE_COST_HEADER] == 0.04 + + def test_litellm_owned_additional_headers_win_over_provider_headers(self): + response = TranscriptionResponse(text="hi") + response._hidden_params = {"additional_headers": {"llm_provider-x-request-id": "kept"}} + + set_provider_response_headers_in_hidden_params(response, {"x-request-id": "provider"}) + + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "kept" + assert response._hidden_params["headers"] == {"x-request-id": "provider"} + + def test_getter_returns_the_recorded_headers(self): + response = ImageResponse() + + set_provider_response_headers_in_hidden_params(response, {"x-request-id": "req_img"}) + + assert get_provider_response_headers_from_hidden_params(response) == {"x-request-id": "req_img"} + + @pytest.mark.parametrize( + "hidden_params", + [ + None, + "headers", + {"additional_headers": {}}, + {"headers": "x-request-id: req_img"}, + {"headers": {"x-request-id": 7}}, + ], + ) + def test_getter_returns_none_without_a_string_header_mapping(self, hidden_params): + response = ImageResponse() + response._hidden_params = hidden_params + + assert get_provider_response_headers_from_hidden_params(response) is None + + def test_getter_returns_none_for_an_object_without_hidden_params(self): + assert get_provider_response_headers_from_hidden_params(object()) is None + + def test_headers_never_leak_into_a_sibling_response(self): + recorded = TranscriptionResponse() + sibling = TranscriptionResponse() + + set_provider_response_headers_in_hidden_params(recorded, {"x-request-id": "req_stt"}) + + assert get_provider_response_headers_from_hidden_params(sibling) is None + assert "additional_headers" not in sibling._hidden_params 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 23c01841b1b..bb8098e6e23 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -24,6 +24,7 @@ from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + _extract_response_obj_and_hidden_params, _get_status_fields, set_callbacks, ) @@ -32,6 +33,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, + ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelResponse, TextCompletionResponse, @@ -8694,3 +8696,88 @@ async def test_async_failure_handler_delivers_failure_payload_to_custom_logger() assert "smoke-failure" in payload["error_str"] assert payload["model"] == "openai/gpt-5.6" assert events.empty() + + +def _image_logging_obj() -> LitellmLogging: + logging_obj = LitellmLogging( + model="gpt-image-2", + messages="a cat", + stream=False, + call_type="aimage_generation", + start_time=time.time(), + litellm_call_id="response-headers-test", + function_id="response-headers-test", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + logging_obj.optional_params = {} + return logging_obj + + +def _image_result_with_headers(request_id: str) -> ImageResponse: + result = ImageResponse(created=1, data=[]) + result._hidden_params = {"headers": {"x-request-id": request_id}} + return result + + +def test_process_hidden_params_surfaces_response_headers_from_the_result(): + logging_obj = _image_logging_obj() + + logging_obj._process_hidden_params_and_response_cost( + _image_result_with_headers("req_img"), datetime.datetime.now(), datetime.datetime.now() + ) + + assert logging_obj.model_call_details["response_headers"] == {"x-request-id": "req_img"} + + +def test_process_hidden_params_keeps_handler_set_response_headers(): + logging_obj = _image_logging_obj() + logging_obj.model_call_details["response_headers"] = {"x-request-id": "from-handler"} + + logging_obj._process_hidden_params_and_response_cost( + _image_result_with_headers("from-result"), datetime.datetime.now(), datetime.datetime.now() + ) + + assert logging_obj.model_call_details["response_headers"] == {"x-request-id": "from-handler"} + + +def _assembled_stream_result_with_headers() -> ModelResponse: + result = _assembled_stream_result() + result._hidden_params = {"headers": {"x-request-id": "req_stream"}} + return result + + +@pytest.mark.asyncio +async def test_async_streaming_success_passes_result_headers_to_callback_kwargs(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result_with_headers()) + + kwargs = releasing.async_log_success_event.await_args.kwargs["kwargs"] + assert kwargs["response_headers"] == {"x-request-id": "req_stream"} + + +def test_sync_streaming_success_passes_result_headers_to_callback_kwargs(): + releasing = CustomLogger() + releasing.log_success_event = MagicMock() + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + + with patcher: + logging_obj.success_handler(result=_assembled_stream_result_with_headers()) + + kwargs = releasing.log_success_event.call_args.kwargs["kwargs"] + assert kwargs["response_headers"] == {"x-request-id": "req_stream"} + + +def test_extract_response_obj_and_hidden_params_reads_binary_content_hidden_params(): + from litellm.types.llms.openai import HttpxBinaryResponseContent as LiteLLMBinaryResponseContent + + result = LiteLLMBinaryResponseContent(response=httpx.Response(status_code=200, content=b"audio bytes")) + result._hidden_params = {"headers": {"x-request-id": "req_tts"}} + + response_obj, hidden_params = _extract_response_obj_and_hidden_params(result, None) + + assert hidden_params == {"headers": {"x-request-id": "req_tts"}} + assert response_obj["object"] == "binary" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 68f37c8ffcc..75b6ce4c626 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -26,6 +26,8 @@ from litellm.llms.base_llm.search.transformation import BaseSearchConfig, Search from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.base_llm.image_generation.transformation import BaseImageGenerationConfig +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -41,7 +43,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @@ -4166,3 +4168,202 @@ async def test_chat_completion_agentic_followup_does_not_repeat_request_params_f assert followup_calls[0]["temperature"] == 0.2 assert followup_calls[0]["api_base"] == "https://a" assert followup_calls[0]["model"] == "openai/gpt-5" + + +_UPSTREAM_HEADERS: Final = {"x-request-id": "req_upstream", "x-ratelimit-remaining-requests": "41"} + + +def _assert_upstream_headers_recorded(response) -> None: + assert response._hidden_params["headers"]["x-request-id"] == "req_upstream" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_upstream" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + + +def _json_with_upstream_headers(payload: dict) -> httpx.MockTransport: + return httpx.MockTransport(lambda request: httpx.Response(200, json=payload, headers=_UPSTREAM_HEADERS)) + + +def _binary_with_upstream_headers() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, content=b"audio-bytes", headers={**_UPSTREAM_HEADERS, "content-type": "audio/mpeg"} + ) + ) + + +def test_audio_transcriptions_records_upstream_response_headers(): + client = HTTPHandler(client=httpx.Client(transport=_json_with_upstream_headers({"text": "transcribed"}))) + + response = BaseLLMHTTPHandler().audio_transcriptions( + client=client, + atranscription=False, + **_json_transcription_call_kwargs(_JSONBodyAudioTranscriptionConfig()), + ) + + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_audio_transcriptions_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"text": "transcribed"})) + + response = await BaseLLMHTTPHandler().async_audio_transcriptions( + client=client, + **_json_transcription_call_kwargs(_JSONBodyAudioTranscriptionConfig()), + ) + + _assert_upstream_headers_recorded(response) + + +def _image_edit_call_kwargs() -> dict: + return { + "model": "edit-model", + "image": b"raw-image", + "prompt": "add a hat", + "image_edit_provider_config": _ImageEditRecordingConfig(), + "image_edit_optional_request_params": {}, + "custom_llm_provider": "openai", + "litellm_params": GenericLiteLLMParams(), + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_image_edit_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_json_with_upstream_headers({"transformed_by": "sync"})) + + response = BaseLLMHTTPHandler().image_edit_handler(client=client, **_image_edit_call_kwargs()) + + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_image_edit_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"transformed_by": "async"})) + + response = await BaseLLMHTTPHandler().async_image_edit_handler(client=client, **_image_edit_call_kwargs()) + + _assert_upstream_headers_recorded(response) + + +class _HeaderImageGenerationConfig(BaseImageGenerationConfig): + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None): + return "https://images.example/v1/generations" + + def transform_image_generation_request(self, model, prompt, optional_params, litellm_params, headers): + return {"prompt": prompt} + + def transform_image_generation_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["b64_json"])]) + + +def _image_generation_call_kwargs() -> dict: + return { + "model": "image-model", + "prompt": "a cat", + "image_generation_provider_config": _HeaderImageGenerationConfig(), + "image_generation_optional_request_params": {}, + "custom_llm_provider": "openai", + "litellm_params": {}, + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_image_generation_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_json_with_upstream_headers({"b64_json": "abc"})) + + response = BaseLLMHTTPHandler().image_generation_handler(client=client, **_image_generation_call_kwargs()) + + assert response.data[0].b64_json == "abc" + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_image_generation_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"b64_json": "abc"})) + + response = await BaseLLMHTTPHandler().async_image_generation_handler( + client=client, **_image_generation_call_kwargs() + ) + + assert response.data[0].b64_json == "abc" + _assert_upstream_headers_recorded(response) + + +class _HeaderTextToSpeechConfig(BaseTextToSpeechConfig): + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, model, optional_params, voice=None, drop_params=False, kwargs=None): + return voice, optional_params + + def validate_environment(self, headers, model, api_key=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://tts.example/v1/speech" + + def transform_text_to_speech_request(self, model, input, voice, optional_params, litellm_params, headers): + return {"dict_body": {"input": input}} + + def transform_text_to_speech_response(self, model, raw_response, logging_obj): + return HttpxBinaryResponseContent(response=raw_response) + + +def _text_to_speech_call_kwargs() -> dict: + return { + "model": "tts-model", + "input": "hello", + "voice": "alloy", + "text_to_speech_provider_config": _HeaderTextToSpeechConfig(), + "text_to_speech_optional_params": {}, + "custom_llm_provider": "openai", + "litellm_params": {}, + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_text_to_speech_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_binary_with_upstream_headers()) + + response = BaseLLMHTTPHandler().text_to_speech_handler(client=client, **_text_to_speech_call_kwargs()) + + assert response.content == b"audio-bytes" + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_text_to_speech_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_binary_with_upstream_headers()) + + response = await BaseLLMHTTPHandler().async_text_to_speech_handler(client=client, **_text_to_speech_call_kwargs()) + + assert response.content == b"audio-bytes" + _assert_upstream_headers_recorded(response) diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py index 9539e13a802..2be691fa65b 100644 --- a/tests/test_litellm/llms/openai/test_openai.py +++ b/tests/test_litellm/llms/openai/test_openai.py @@ -1,13 +1,15 @@ import asyncio import json from typing import Final +from unittest.mock import Mock import httpx import pytest -from openai import AsyncOpenAI +from openai import AsyncOpenAI, OpenAI import litellm from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.types.utils import ImageResponse @pytest.mark.parametrize( @@ -253,3 +255,98 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport() assert tool_call.function.name == "get_weather" assert json.loads(tool_call.function.arguments) == {"city": "Paris"} assert rebuilt.choices[0].finish_reason == "tool_calls" + + +_PROVIDER_HEADERS: Final = {"x-request-id": "req_openai", "x-ratelimit-remaining-requests": "41"} + + +def _image_generation_transport() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"created": 1, "data": [{"b64_json": "abc"}]}, headers=_PROVIDER_HEADERS + ) + ) + + +def _speech_transport() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, content=b"audio-bytes", headers={**_PROVIDER_HEADERS, "content-type": "audio/mpeg"} + ) + ) + + +def _assert_provider_headers_recorded(response) -> None: + assert response._hidden_params["headers"]["x-request-id"] == "req_openai" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_openai" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + + +def _image_generation_kwargs() -> dict: + return { + "model": "gpt-image-2", + "prompt": "a cat", + "timeout": 10, + "optional_params": {}, + "logging_obj": Mock(), + "api_key": "transport-only", + "model_response": ImageResponse(), + } + + +def test_image_generation_records_provider_response_headers(): + with httpx.Client(transport=_image_generation_transport()) as http_client: + response = OpenAIChatCompletion().image_generation( + client=OpenAI(api_key="transport-only", http_client=http_client), **_image_generation_kwargs() + ) + + _assert_provider_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_aimage_generation_records_provider_response_headers(): + async with httpx.AsyncClient(transport=_image_generation_transport()) as http_client: + response = await OpenAIChatCompletion().image_generation( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + aimg_generation=True, + **_image_generation_kwargs(), + ) + + _assert_provider_headers_recorded(response) + + +def _audio_speech_kwargs() -> dict: + return { + "model": "gpt-4o-mini-tts", + "input": "hello", + "voice": "alloy", + "optional_params": {}, + "api_key": "transport-only", + "api_base": None, + "organization": None, + "project": None, + "max_retries": 0, + "timeout": 10, + "logging_obj": Mock(), + } + + +def test_audio_speech_records_provider_response_headers(): + with httpx.Client(transport=_speech_transport()) as http_client: + response = OpenAIChatCompletion().audio_speech( + client=OpenAI(api_key="transport-only", http_client=http_client), **_audio_speech_kwargs() + ) + + _assert_provider_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_audio_speech_records_provider_response_headers(): + async with httpx.AsyncClient(transport=_speech_transport()) as http_client: + response = await OpenAIChatCompletion().audio_speech( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + aspeech=True, + **_audio_speech_kwargs(), + ) + + _assert_provider_headers_recorded(response) diff --git a/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py b/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py new file mode 100644 index 00000000000..f2dbb71fea1 --- /dev/null +++ b/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py @@ -0,0 +1,71 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +from openai import AsyncOpenAI, OpenAI + +from litellm.llms.openai.transcriptions.handler import OpenAIAudioTranscription +from litellm.types.utils import TranscriptionResponse + +_PROVIDER_HEADERS: Final = {"x-request-id": "req_stt", "x-ratelimit-remaining-requests": "41"} + + +def _transcription_transport() -> httpx.MockTransport: + return httpx.MockTransport(lambda request: httpx.Response(200, json={"text": "hello"}, headers=_PROVIDER_HEADERS)) + + +def _logging_obj() -> Mock: + logging_obj = Mock() + logging_obj.model_call_details = {} + return logging_obj + + +def _call_kwargs(logging_obj: Mock) -> dict: + return { + "model": "gpt-4o-mini-transcribe", + "audio_file": ("audio.wav", b"riff-bytes", "audio/wav"), + "optional_params": {}, + "litellm_params": {}, + "model_response": TranscriptionResponse(), + "timeout": 10.0, + "max_retries": 0, + "logging_obj": logging_obj, + "api_key": "transport-only", + "api_base": None, + } + + +def _assert_headers_recorded(response: TranscriptionResponse, logging_obj: Mock) -> None: + assert response.text == "hello" + assert response._hidden_params["headers"]["x-request-id"] == "req_stt" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_stt" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + assert logging_obj.model_call_details["response_headers"]["x-request-id"] == "req_stt" + + +def test_audio_transcriptions_records_provider_response_headers(): + logging_obj = _logging_obj() + + with httpx.Client(transport=_transcription_transport()) as http_client: + response = OpenAIAudioTranscription().audio_transcriptions( + client=OpenAI(api_key="transport-only", http_client=http_client), + atranscription=False, + **_call_kwargs(logging_obj), + ) + + _assert_headers_recorded(response, logging_obj) + + +@pytest.mark.asyncio +async def test_async_audio_transcriptions_records_provider_response_headers(): + logging_obj = _logging_obj() + + async with httpx.AsyncClient(transport=_transcription_transport()) as http_client: + response = await OpenAIAudioTranscription().audio_transcriptions( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + atranscription=True, + **_call_kwargs(logging_obj), + ) + + _assert_headers_recorded(response, logging_obj) diff --git a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py b/tests/test_litellm/test_non_chat_routes_open_llm_spans.py index d62959ccd43..02c95c4bb2a 100644 --- a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py +++ b/tests/test_litellm/test_non_chat_routes_open_llm_spans.py @@ -46,9 +46,9 @@ class _FakeSpeech: )() -class _FakeImages: +class _FakeRawImages: async def generate(self, **kwargs: Any) -> Any: - return type( + parsed: Final = type( "_Images", (), { @@ -58,6 +58,16 @@ class _FakeImages: } }, )() + return type( + "_RawImages", + (), + {"parse": lambda self: parsed, "headers": httpx.Headers({"x-request-id": "req-image"})}, + )() + + +class _FakeImages: + def __init__(self) -> None: + self.with_raw_response = _FakeRawImages() class _FakeModerations: diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 55ef74abd7b..11df07f5fea 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -12,6 +12,14 @@ import pytest from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.types.utils import ImageResponse + + +def _raw_image_response(mock_image_data): + raw_response = MagicMock() + raw_response.parse.return_value = mock_image_data + raw_response.headers = {"x-request-id": "req-image"} + return raw_response @pytest.fixture @@ -41,7 +49,7 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -58,7 +66,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers def test_sync_image_generation_without_headers( @@ -72,7 +80,7 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -86,7 +94,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert "extra_headers" not in kwargs @pytest.mark.asyncio @@ -101,7 +109,9 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" test_headers = {"cf-aig-authorization": "Bearer custom-token"} @@ -109,7 +119,7 @@ class TestImageGenerationExtraHeaders: await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", @@ -117,7 +127,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers @pytest.mark.asyncio @@ -132,20 +142,22 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert "extra_headers" not in kwargs @pytest.mark.parametrize("is_async", [False, True]) @@ -169,11 +181,13 @@ class TestImageGenerationExtraHeaders: test_headers = {"cf-aig-authorization": "Bearer custom-token"} if is_async: - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", @@ -181,7 +195,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) else: - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) openai_chat_completions.image_generation( model="dall-e-3", prompt="A white cat", @@ -197,7 +211,7 @@ class TestImageGenerationExtraHeaders: "complete_input_dict" ] assert "extra_headers" not in logged_body - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers def test_sync_image_generation_forwards_headers_to_async( @@ -242,7 +256,9 @@ class TestImageGenerationEntryPointHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -256,6 +272,6 @@ class TestImageGenerationEntryPointHeaders: api_key="test-key", ) - mock_openai_client.images.generate.assert_called_once() - _, kwargs = mock_openai_client.images.generate.call_args + mock_openai_client.images.with_raw_response.generate.assert_called_once() + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index b5eec42b569..ee7bdebe745 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -526,6 +526,7 @@ class TestVertexAILyriaTextToSpeechConfig: ): mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} mock_response.json.return_value = response_json with ( patch.object( # test-quality-ok: litellm.speech has no seam for Vertex token minting From 5a8ec1378611e6620f5cfe291f0e971ca403293e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:04:03 -0500 Subject: [PATCH 004/154] feat(usage): search keys beyond the top-N usage subset (#42827) --- litellm/proxy/_types.py | 2 + .../internal_user_endpoints.py | 153 ++++++++++++++-- .../internal_user_endpoints.py | 12 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../proxy/auth/test_route_checks.py | 1 + .../test_internal_user_endpoints.py | 166 ++++++++++++++++++ .../_components/components/UsagePageView.tsx | 16 +- .../components/KeyActivityPanel.test.tsx | 64 +++++++ .../UsagePage/components/KeyActivityPanel.tsx | 73 +++++++- .../src/components/networking.tsx | 30 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 65 +++++++ 11 files changed, 560 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4affa55f903..51a3eb03067 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -699,6 +699,7 @@ class LiteLLMRoutes(enum.Enum): "/user/list", "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", # team "/team/new", "/team/update", @@ -901,6 +902,7 @@ class LiteLLMRoutes(enum.Enum): "/model/delete", "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", # Endpoint restricts results to organizations the caller is ORG_ADMIN # of; a caller who administers none gets an empty result set. "/organization/daily/activity", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 59d8dd821d8..7b25348aa53 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -28,6 +28,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( @@ -87,11 +88,13 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, + KeyActivitySearchWhere, UserListResponse, UserSearchWhere, UserUpdateResult, @@ -2991,6 +2994,27 @@ async def get_user_daily_activity( ) +def _resolve_user_daily_activity_entity_id( + user_api_key_dict: UserAPIKeyAuth, + user_id: str | None, +) -> str | None: + is_admin: Final = _user_has_admin_view(user_api_key_dict) + + if is_admin: + return user_id + + caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) + effective_user_id: Final = user_id if user_id is not None else caller_user_id + if effective_user_id != caller_user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI detail payload shape + "error": "Non-admin users can only view their own spend data." + }, + ) + return effective_user_id + + @router.get( "/user/daily/activity/aggregated", tags=["Budget & Spend Tracking", "Internal User management"], @@ -3057,20 +3081,7 @@ async def get_user_daily_activity_aggregated( ) try: - is_admin: Final = _user_has_admin_view(user_api_key_dict) - - if is_admin: - entity_id = user_id # None means global view, otherwise filter by user - else: - caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) - if user_id is None: - user_id = caller_user_id - if user_id != caller_user_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Non-admin users can only view their own spend data."}, - ) - entity_id = user_id + entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) return await get_daily_activity_aggregated( prisma_client=prisma_client, @@ -3094,3 +3105,117 @@ async def get_user_daily_activity_aggregated( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {e}"}, ) + + +@router.get( + "/user/daily/activity/aggregated/search", + tags=["Budget & Spend Tracking", "Internal User management"], # mutable-ok: FastAPI route tags shape + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route dependencies shape + response_model=SpendAnalyticsPaginatedResponse, +) +@management_endpoint_wrapper +async def search_user_daily_activity_keys( + search: str = fastapi.Query( + ..., + min_length=1, + description="Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive)", + ), + start_date: str | None = fastapi.Query( + default=None, + description="Start date in YYYY-MM-DD format", + ), + end_date: str | None = fastapi.Query( + default=None, + description="End date in YYYY-MM-DD format", + ), + user_id: str | None = fastapi.Query( + default=None, + description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", + ), + timezone: int | None = fastapi.Query( + default=None, + description="Timezone offset in minutes from UTC (e.g., 480 for PST). " + "Matches JavaScript's Date.getTimezoneOffset() convention.", + ), + include_current_utc_day: bool = fastapi.Query( + default=False, + description="When the range ends on the caller's current local day, extend it to " + "today's UTC bucket so spend written after the caller's local midnight (in UTC " + "terms) is included. Requires the timezone parameter. Historical ranges are " + "never extended.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> SpendAnalyticsPaginatedResponse: + """ + Search verification tokens by exact token hash or by a case-insensitive substring of + the key alias or owning user ID, then return the aggregated daily activity for the + matches. Lets the Usage page surface keys that fell outside the top-spend subset + the aggregated endpoint loads. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: FastAPI detail payload shape + "error": CommonProxyErrors.db_not_connected_error.value + }, + ) + + if start_date is None or end_date is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Please provide start_date and end_date"}, # mutable-ok: FastAPI detail payload shape + ) + + try: + entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) + + search_or: Final = ( + {"token": search}, # mutable-ok: prisma serializes where clauses, keep plain dicts + {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + ) + where: Final[KeyActivitySearchWhere] = ( + {"OR": search_or} # mutable-ok: prisma where clause root + if entity_id is None + else {"user_id": entity_id, "OR": search_or} # mutable-ok: prisma where clause root + ) + matched_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many( + where=where, + take=USAGE_TOP_API_KEYS_LIMIT, + order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict + ) + tokens: Final = [key.token for key in matched_keys] # mutable-ok: api_key filter union expects a list + + if not tokens: + return SpendAnalyticsPaginatedResponse( + results=[], # mutable-ok: response model field shape + metadata=DailySpendMetadata( + api_key_limit=USAGE_TOP_API_KEYS_LIMIT, + total_api_keys=0, + ), + ) + + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=tokens, + timezone_offset_minutes=timezone, + include_current_utc_day=include_current_utc_day, + ) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("/user/daily/activity/aggregated/search: Exception occured - %s", e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Failed to fetch analytics: {e}"}, # mutable-ok: FastAPI detail payload shape + ) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 43e3899d523..05cbd4507a2 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, @@ -28,6 +28,16 @@ class UserSearchWhere(TypedDict): OR: ReadOnly[tuple[Mapping[Literal["user_id", "user_email"], InsensitiveContains], ...]] +class KeyActivitySearchWhere(TypedDict): + """Prisma filter behind `/user/daily/activity/aggregated/search`: exact token hash, or key alias + or user id containing the term, case-insensitive.""" + + user_id: NotRequired[ReadOnly[str]] + OR: ReadOnly[ + tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] + ] + + class UserListResponse(BaseModel): """ Response model for the user list endpoint diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index f8277c83a64..d0f9c31ecaf 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -32,6 +32,7 @@ GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity GET /user/daily/activity/aggregated +GET /user/daily/activity/aggregated/search GET /user/spend/report # Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index e5179387f82..571e066e947 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3517,6 +3517,7 @@ def test_internal_user_still_blocked_from_another_users_info(): [ "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", ], ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index c663e63414c..8260aec9326 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2659,6 +2659,172 @@ async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_us assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123" +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_passes_matched_tokens_to_aggregation(monkeypatch): + """The search endpoint resolves matching verification tokens by hash, alias, or + user id, then aggregates daily spend for exactly those tokens. This is what lets + the Usage page find keys outside the top-spend subset the aggregated endpoint caps.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token="tok-a"), SimpleNamespace(token="tok-b")] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_response = MagicMock() + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + admin_key_dict = UserAPIKeyAuth( + user_id="admin-user-001", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + result = await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=480, + include_current_utc_day=False, + user_api_key_dict=admin_key_dict, + ) + + assert result is mock_response + + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert find_many_kwargs["where"]["OR"] == ( + {"token": "gamma"}, + {"key_alias": {"contains": "gamma", "mode": "insensitive"}}, + {"user_id": {"contains": "gamma", "mode": "insensitive"}}, + ) + assert "user_id" not in find_many_kwargs["where"] + + mock_get_daily_agg.assert_called_once_with( + prisma_client=mock_prisma_client, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2025-02-01", + end_date="2025-02-28", + model=None, + api_key=["tok-a", "tok-b"], + timezone_offset_minutes=480, + include_current_utc_day=False, + ) + + +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_no_match_returns_empty_without_aggregating(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_get_daily_agg = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + admin_key_dict = UserAPIKeyAuth( + user_id="admin-user-001", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + result = await search_user_daily_activity_keys( + search="nothing-matches", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=admin_key_dict, + ) + + assert result.results == [] + assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT + assert result.metadata.total_api_keys == 0 + mock_get_daily_agg.assert_not_called() + + +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_non_admin_scoped_to_caller(monkeypatch): + """Same scoping contract as the aggregated route: a non-admin with no user_id + is scoped to their own rows, and any other user_id is a 403.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[SimpleNamespace(token="tok-a")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + non_admin_key_dict = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_response = MagicMock() + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + result = await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=non_admin_key_dict, + ) + + assert result is mock_response + assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "user-1" + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs["where"]["user_id"] == "user-1" + + with pytest.raises(HTTPException) as exc_info: + await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id="user-2", + timezone=None, + include_current_utc_day=False, + user_api_key_dict=non_admin_key_dict, + ) + + assert exc_info.value.status_code == 403 + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_delete_user_cleans_up_created_by_invitation_links(mocker): """ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 228d8acf146..e0171d97423 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -39,6 +39,7 @@ import { tagListCall, userDailyActivityAggregatedCall, userDailyActivityCall, + userDailyActivityKeySearchCall, } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { ChartLoader } from "@/components/shared/chart_loader"; @@ -437,6 +438,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { [userSpendData, modelViewType, teams], ); const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]); + const searchKeys = useCallback( + (q: string) => { + if (!accessToken || !startTime || !endTime) return Promise.resolve({}); + return userDailyActivityKeySearchCall(accessToken, startTime, endTime, q, effectiveUserId).then((data) => + processActivityData(data, "api_keys", teams), + ); + }, + [accessToken, startTime, endTime, effectiveUserId, teams], + ); const mcpServerMetrics = useMemo( () => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams], @@ -865,7 +875,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx index 830139143e9..f8a7d07633b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -78,4 +78,68 @@ describe("KeyActivityPanel", () => { render(); expect(screen.queryByRole("note")).not.toBeInTheDocument(); }); + + it("finds keys outside the loaded top-spend subset via server search", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + + expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); + expect(searchKeys).toHaveBeenCalledWith("gamma"); + expect(screen.getByText("Showing 1 of 3 keys")).toBeInTheDocument(); + }); + + it("never calls the server search when every key is already loaded", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + render(); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + + expect(await screen.findByText('No keys match "gamma" in this date range')).toBeInTheDocument(); + await new Promise((resolve) => setTimeout(resolve, 400)); + expect(searchKeys).not.toHaveBeenCalled(); + }); + + it("drops stale server results as soon as the search callback is rebuilt", async () => { + const searchKeysA = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + const searchKeysB = vi + .fn<(query: string) => Promise>>() + .mockReturnValue(new Promise(() => {})); + const { rerender } = render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); + + rerender( + , + ); + + expect(screen.getByRole("status")).toHaveTextContent("Searching all keys"); + expect(screen.queryByText("hash-gamma")).not.toBeInTheDocument(); + }); + + it("reports a failed server search but keeps the local matches", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockRejectedValue(new Error("boom")); + render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "alice" } }); + + expect(await screen.findByRole("alert")).toHaveTextContent("Key search failed"); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice"); + }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx index 8b2141f8528..3467ba61b22 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -1,5 +1,5 @@ import { Search, X } from "lucide-react"; -import React, { useMemo, useState } from "react"; +import React, { useEffect, useMemo, useState } from "react"; import { ActivityMetrics } from "@/components/activity_metrics"; import type { ApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; @@ -12,18 +12,67 @@ interface KeyActivityPanelProps { keyMetrics: Record; hidePromptCachingMetrics?: boolean; apiKeyTruncation?: ApiKeyTruncation; + searchKeys?: SearchKeys; } +type SearchKeys = (query: string) => Promise>; + +type RemoteSearch = + | { status: "idle" } + | { status: "loading"; query: string; searchKeys: SearchKeys } + | { status: "done"; query: string; searchKeys: SearchKeys; keys: Record } + | { status: "error"; query: string; searchKeys: SearchKeys }; + +const REMOTE_SEARCH_DEBOUNCE_MS = 300; + const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false, apiKeyTruncation, + searchKeys, }) => { const [query, setQuery] = useState(""); + const [remote, setRemote] = useState({ status: "idle" }); const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); + const trimmedQuery = query.trim(); + const remoteEnabled = searchKeys !== undefined && apiKeyTruncation !== undefined && trimmedQuery !== ""; + + useEffect(() => { + if (!remoteEnabled) return; + let cancelled = false; + const timer = setTimeout(() => { + setRemote({ status: "loading", query: trimmedQuery, searchKeys }); + searchKeys(trimmedQuery) + .then((keys) => { + if (!cancelled) setRemote({ status: "done", query: trimmedQuery, searchKeys, keys }); + }) + .catch(() => { + if (!cancelled) setRemote({ status: "error", query: trimmedQuery, searchKeys }); + }); + }, REMOTE_SEARCH_DEBOUNCE_MS); + return () => { + cancelled = true; + clearTimeout(timer); + }; + }, [remoteEnabled, trimmedQuery, searchKeys]); + + const remoteMatchesSearch = + "searchKeys" in remote && remote.searchKeys === searchKeys && remote.query === trimmedQuery; + const remoteCurrent = remoteEnabled && remoteMatchesSearch; + const remoteLoading = remoteEnabled && (remote.status === "loading" || !remoteCurrent); + const remoteFailed = remoteCurrent && remote.status === "error"; + + const extraRemoteKeys = useMemo(() => { + const remoteKeys = remoteCurrent && remote.status === "done" ? remote.keys : {}; + return Object.fromEntries(Object.entries(remoteKeys).filter(([hash]) => !(hash in keyMetrics))); + }, [remoteCurrent, remote, keyMetrics]); + const displayed = useMemo(() => ({ ...extraRemoteKeys, ...filtered }), [extraRemoteKeys, filtered]); + const totalKeys = Object.keys(keyMetrics).length; - const shownKeys = Object.keys(filtered).length; - const isFiltering = query.trim() !== ""; + const shownKeys = Object.keys(displayed).length; + const totalShown = totalKeys + Object.keys(extraRemoteKeys).length; + const isFiltering = trimmedQuery !== ""; + const noMatches = isFiltering && !remoteLoading && totalKeys > 0 && shownKeys === 0; return (
@@ -47,8 +96,18 @@ const KeyActivityPanel: React.FC = ({ )} - Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys + Showing {shownKeys.toLocaleString()} of {totalShown.toLocaleString()} keys + {remoteLoading && ( + + Searching all keys... + + )} + {remoteFailed && ( + + Key search failed + + )} {apiKeyTruncation !== undefined && ( Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} @@ -56,12 +115,12 @@ const KeyActivityPanel: React.FC = ({ )}
- {isFiltering && totalKeys > 0 && shownKeys === 0 ? ( + {noMatches ? (

- No keys match "{query.trim()}" in this date range + No keys match "{trimmedQuery}" in this date range

) : ( - + )} ); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index c2f4fa80634..e6923a26d4c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2556,6 +2556,36 @@ export const userDailyActivityAggregatedCall = async ( } }; +export const userDailyActivityKeySearchCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + ...options: [search: string, userId?: string | null] +) => { + const [search, userId = null] = options; + try { + const formatDate = (date: Date) => { + const year = date.getFullYear(); + const month = String(date.getMonth() + 1).padStart(2, "0"); + const day = String(date.getDate()).padStart(2, "0"); + return `${year}-${month}-${day}`; + }; + return await apiClient.get(`/user/daily/activity/aggregated/search`, { + accessToken, + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + search, + user_id: userId || undefined, + }, + }); + } catch (error) { + console.error("Failed to search user daily activity keys:", error); + throw error; + } +}; + export const gatewayDailyActivityCall = async (accessToken: string, startTime: Date, endTime: Date) => { /** * Get gateway request counts (SGR) recorded by the proxy middleware. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 052cc693f8b..79d008c1b48 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -17326,6 +17326,29 @@ export interface paths { patch?: never; trace?: never; }; + "/user/daily/activity/aggregated/search": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Search User Daily Activity Keys + * @description Search verification tokens by exact token hash or by a case-insensitive substring of + * the key alias or owning user ID, then return the aggregated daily activity for the + * matches. Lets the Usage page surface keys that fell outside the top-spend subset + * the aggregated endpoint loads. + */ + get: operations["search_user_daily_activity_keys_user_daily_activity_aggregated_search_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/user/delete": { parameters: { query?: never; @@ -68731,6 +68754,48 @@ export interface operations { }; }; }; + search_user_daily_activity_keys_user_daily_activity_aggregated_search_get: { + parameters: { + query: { + /** @description Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive) */ + search: string; + /** @description Start date in YYYY-MM-DD format */ + start_date?: string | null; + /** @description End date in YYYY-MM-DD format */ + end_date?: string | null; + /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ + user_id?: string | null; + /** @description Timezone offset in minutes from UTC (e.g., 480 for PST). Matches JavaScript's Date.getTimezoneOffset() convention. */ + timezone?: number | null; + /** @description When the range ends on the caller's current local day, extend it to today's UTC bucket so spend written after the caller's local midnight (in UTC terms) is included. Requires the timezone parameter. Historical ranges are never extended. */ + include_current_utc_day?: boolean; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_user_user_delete_post: { parameters: { query?: never; From c2eb549ee685b442b74bc7c1cf9c4c1a7698a5ec Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:04:12 -0500 Subject: [PATCH 005/154] feat(usage): search team keys beyond the top-N in the Team usage view (#42857) --- litellm/proxy/_types.py | 5 + .../management_endpoints/team_endpoints.py | 109 +++++++++ .../management_endpoints/team_endpoints.py | 24 ++ .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_daily_activity_key_search.py | 128 +++++++++++ .../management/test_team_daily_activity.py | 15 +- .../proxy/auth/test_route_checks.py | 49 ++++ .../test_team_endpoints.py | 212 ++++++++++++++++++ .../components/EntityUsage/EntityUsage.tsx | 14 +- .../src/components/networking.tsx | 25 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 60 ++++- 11 files changed, 635 insertions(+), 7 deletions(-) create mode 100644 tests/integration/spend/test_team_daily_activity_key_search.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 51a3eb03067..73e3e0e6ee0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -307,6 +307,7 @@ class KeyManagementRoutes(str, enum.Enum): # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated" + TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search" # team spend-log viewing SPEND_LOGS = "/spend/logs" @@ -673,6 +674,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value, + KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH.value, KeyManagementRoutes.SPEND_LOGS.value, KeyManagementRoutes.SPEND_LOGS_V2.value, KeyManagementRoutes.KEY_RESET_SPEND.value, @@ -717,6 +719,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only "/gateway/daily/activity", @@ -887,6 +890,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/team/spend/by_user", "/team/{team_id}/members/me", # POST/GET the team's logging callbacks, and DELETE one of them. Every @@ -986,6 +990,7 @@ class LiteLLMRoutes(enum.Enum): "/user/daily/activity", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/tag/daily/activity", "/tag/list", "/audit", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6cec3e714ec..493c83c730a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -40,6 +40,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -196,6 +197,7 @@ from litellm.repositories.verification_token_repository import ( from litellm.router import Router from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.team_endpoints import ( @@ -204,7 +206,9 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, + TeamIdSearchFilter, TeamIdSearchMatch, + TeamKeyActivitySearchWhere, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -6805,6 +6809,111 @@ async def get_team_daily_activity_aggregated( ) +def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere: + """Caller scoping lives inside the same Prisma where as the search term so `take` + never trims visible matches in favour of keys the caller is not allowed to see.""" + search_or: Final = ( + {"token": search}, # mutable-ok: prisma where clause leaf + {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + ) + own_keys: Final = tuple(scope.api_key_filter) if isinstance(scope.api_key_filter, list) else None + team_filter: Final[TeamIdSearchFilter | None] = ( + { # mutable-ok: prisma where clause leaf + "in": tuple(scope.team_ids), + "notIn": tuple(scope.exclude_team_ids), + } + if scope.team_ids is not None and scope.exclude_team_ids is not None + else {"in": tuple(scope.team_ids)} # mutable-ok: prisma where clause leaf + if scope.team_ids is not None + else {"notIn": tuple(scope.exclude_team_ids)} # mutable-ok: prisma where clause leaf + if scope.exclude_team_ids is not None + else None + ) + if team_filter is None and own_keys is None: + return {"OR": search_or} # mutable-ok: prisma where clause root + if team_filter is None and own_keys is not None: + return {"token": {"in": own_keys}, "OR": search_or} # mutable-ok: prisma where clause root + if team_filter is not None and own_keys is None: + return {"team_id": team_filter, "OR": search_or} # mutable-ok: prisma where clause root + assert team_filter is not None and own_keys is not None + return { # mutable-ok: prisma where clause root + "team_id": team_filter, + "token": {"in": own_keys}, # mutable-ok: prisma where clause leaf + "OR": search_or, + } + + +@router.get( + "/team/daily/activity/aggregated/search", + response_model=SpendAnalyticsPaginatedResponse, + tags=["team management"], # mutable-ok: FastAPI route tags shape +) +async def search_team_daily_activity_keys( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + search: str = fastapi.Query( + ..., + min_length=1, + description="Exact token hash, or a case-insensitive substring of the key alias or owning user id", + ), + team_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + exclude_team_ids: str | None = None, + timezone: int | None = None, +) -> SpendAnalyticsPaginatedResponse: + """Aggregated daily team activity for the keys matching `search`, across every key the caller may + see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend.""" + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) + + range_error: Final = _aggregated_date_range_error(start_date, end_date) + if range_error is not None: + raise _daily_activity_error(status_code=400, message=range_error) + + scope: Final = await _resolve_team_daily_activity_scope( + team_ids=team_ids, + exclude_team_ids=exclude_team_ids, + api_key=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + matched_keys: Final = await _tokens_db(prisma_client).find_many( + where=_team_key_search_where(search=search, scope=scope), + take=USAGE_TOP_API_KEYS_LIMIT, + order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict + ) + tokens: Final = [key.token for key in matched_keys] # mutable-ok: get_daily_activity_aggregated takes list[str] + if not tokens: + return SpendAnalyticsPaginatedResponse( + results=[], # mutable-ok: response model field shape + metadata=DailySpendMetadata(api_key_limit=USAGE_TOP_API_KEYS_LIMIT, total_api_keys=0), + ) + + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=scope.team_ids, + entity_metadata_field=scope.team_alias_metadata, + start_date=start_date, + end_date=end_date, + model=None, + api_key=tokens, + exclude_entity_ids=scope.exclude_team_ids, + timezone_offset_minutes=timezone, + include_entity_breakdown=True, + ) + + def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count)) user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else "" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 4524c47ec38..aac2703e918 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.proxy._types import ( KeyManagementRoutes, @@ -11,10 +13,32 @@ from litellm.proxy._types import ( MemberDeleteRequest, ) from litellm.proxy.common_utils.timezone_utils import budget_duration_error +from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse TeamIdSearchMatch = Literal["exact", "prefix"] + +TeamIdSearchFilter = TypedDict( + "TeamIdSearchFilter", + { # mutable-ok: functional TypedDict field map + "in": NotRequired[ReadOnly[Sequence[str]]], + "notIn": NotRequired[ReadOnly[Sequence[str]]], + }, +) + + +class TeamKeyActivitySearchWhere(TypedDict): + """Prisma filter behind `/team/daily/activity/aggregated/search`: exact token hash, or key alias + or user id containing the term, case-insensitive, narrowed to the teams and keys the caller may see.""" + + team_id: NotRequired[ReadOnly[TeamIdSearchFilter]] + token: NotRequired[ReadOnly[Mapping[Literal["in"], Sequence[str]]]] + OR: ReadOnly[ + tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] + ] + + MAX_BULK_TEAM_MEMBER_DELETES: Final = 500 MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES: Final = 500 diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d0f9c31ecaf..b0f28a9c740 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated +GET /team/daily/activity/aggregated/search GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity diff --git a/tests/integration/spend/test_team_daily_activity_key_search.py b/tests/integration/spend/test_team_daily_activity_key_search.py new file mode 100644 index 00000000000..2b0395cf4ba --- /dev/null +++ b/tests/integration/spend/test_team_daily_activity_key_search.py @@ -0,0 +1,128 @@ +import uuid +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from pydantic import JsonValue + +_SEARCH_PATH: Final = "/team/daily/activity/aggregated/search" + + +def _range_around_today() -> dict[str, str]: + today: Final = datetime.now(timezone.utc) + return { + "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), + "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"), + "timezone": "0", + } + + +def _team_key_breakdown(body: dict[str, JsonValue], team: str) -> dict[str, JsonValue]: + results: Final = body["results"] + assert isinstance(results, list) and len(results) == 1, body + entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) + return object_value(object_value(entities[team])["api_key_breakdown"]) + + +def test_team_key_search_returns_only_the_matching_key_spend_by_alias_and_by_hash(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + needle_alias: Final = f"needle-{uuid.uuid4().hex}" + needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) + other: Final = scenario.key(team_id=team, models=[model], key_alias=f"other-{uuid.uuid4().hex}") + needle_digest: Final = sha256(needle.encode()).hexdigest() + other_digest: Final = sha256(other.encode()).hexdigest() + for key in (needle, other): + reply: Final = gateway.chat(model, key=key, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: sorted(row["api_key"] for row in values) == sorted((needle_digest, other_digest)), + seconds=70, + ) + assert all(float(row["spend"]) == pytest.approx(0.06) for row in daily), daily + for search in (needle_alias.upper(), needle_digest): + response: Final = gateway.request( + "GET", _SEARCH_PATH, params={"team_ids": team, "search": search, **_range_around_today()} + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + assert object_value(body["metadata"])["total_spend"] == pytest.approx(0.06), response.text + per_key: Final = _team_key_breakdown(body, team) + assert set(per_key) == {needle_digest}, response.text + assert object_value(object_value(per_key[needle_digest])["metrics"])["spend"] == pytest.approx(0.06) + + +def test_team_key_search_is_scoped_to_the_teams_the_caller_belongs_to(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + needle_alias: Final = f"needle-{uuid.uuid4().hex}" + needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) + needle_digest: Final = sha256(needle.encode()).hexdigest() + reply: Final = gateway.chat(model, key=needle, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + eventually( + lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: [row["api_key"] for row in values] == [needle_digest], + seconds=70, + ) + outsider: Final = scenario.user(user_role="internal_user") + outsider_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": outsider, "role": "user"}]) + outsider_key: Final = scenario.key(user_id=outsider, team_id=outsider_team, models=[model]) + params: Final = {"search": needle_alias, **_range_around_today()} + admin_view: Final = gateway.request("GET", _SEARCH_PATH, params={"team_ids": team, **params}) + assert admin_view.status_code == 200, admin_view.text + assert set(_team_key_breakdown(object_value(admin_view.json()), team)) == {needle_digest}, admin_view.text + own_teams_view: Final = gateway.request("GET", _SEARCH_PATH, params=params, key=outsider_key) + assert own_teams_view.status_code == 200, own_teams_view.text + own_teams_body: Final = object_value(own_teams_view.json()) + assert own_teams_body["results"] == [], own_teams_view.text + assert object_value(own_teams_body["metadata"])["total_api_keys"] == 0, own_teams_view.text + foreign_team_view: Final = gateway.request( + "GET", _SEARCH_PATH, params={"team_ids": team, **params}, key=outsider_key + ) + assert foreign_team_view.status_code == 404, foreign_team_view.text + + +def test_team_key_search_excludes_teams_inside_the_where(gateway: Gateway) -> None: + """The dashboard always sends exclude_team_ids; a matching key in an excluded + team with higher spend must not consume a take slot nor appear in the result.""" + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_keep: Final = scenario.team(models=[model]) + team_drop: Final = scenario.team(models=[model]) + shared_alias: Final = f"needle-{uuid.uuid4().hex}" + keep: Final = scenario.key(team_id=team_keep, models=[model], key_alias=f"{shared_alias}-keep") + drop: Final = scenario.key(team_id=team_drop, models=[model], key_alias=f"{shared_alias}-drop") + keep_digest: Final = sha256(keep.encode()).hexdigest() + drop_digest: Final = sha256(drop.encode()).hexdigest() + for _ in range(2): + reply: Final = gateway.chat(model, key=drop, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + reply = gateway.chat(model, key=keep, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + eventually( + lambda: read_rows( + 'SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id IN (%s, %s)', + (team_keep, team_drop), + ), + lambda values: sorted(row["api_key"] for row in values) == sorted((keep_digest, drop_digest)), + seconds=70, + ) + response: Final = gateway.request( + "GET", + _SEARCH_PATH, + params={"search": shared_alias, "exclude_team_ids": team_drop, **_range_around_today()}, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + results: Final = body["results"] + assert isinstance(results, list) and len(results) == 1, body + entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) + assert set(entities) == {team_keep}, response.text + assert set(_team_key_breakdown(body, team_keep)) == {keep_digest}, response.text diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index d84cc4c94af..9bbc8fdde29 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -5,8 +5,9 @@ from .actors import Actor pytestmark = pytest.mark.asyncio(loop_scope="session") -# GET /team/daily/activity and its /aggregated variant (same shared scope -# resolver, so the matrix must hold for both). A proxy admin (admin view) sees +# GET /team/daily/activity, its /aggregated variant, and the key-search +# variant (same shared scope resolver, so the matrix must hold for all +# three). A proxy admin (admin view) sees # activity for any team. A non-admin is scoped to user_info.teams: a bare query # defaults to its own teams (200), and an explicit team_ids filter naming a # team it does not belong to is 404 (the VERIA-43 fix). Org admins have no @@ -43,8 +44,12 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" @pytest.mark.parametrize( "endpoint", - ("/team/daily/activity", "/team/daily/activity/aggregated"), - ids=("paginated", "aggregated"), + ( + "/team/daily/activity", + "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", + ), + ids=("paginated", "aggregated", "search"), ) @pytest.mark.parametrize( "actor,team,expected_status", @@ -54,7 +59,7 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" async def test_team_daily_activity_matrix( actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world ): - query = _DATES + query = _DATES + ("&search=x" if endpoint.endswith("/search") else "") if team == "alpha": query += f"&team_ids={world.team_alpha_id}" elif team == "beta": diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 571e066e947..7bb79a115dd 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3600,6 +3600,55 @@ def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): ) +@pytest.mark.parametrize( + "route", + [ + "/team/daily/activity", + "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", + ], +) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_team_daily_activity_routes_reachable_by_non_admin(route, user_role): + """The Team Usage dashboard calls all three team daily-activity routes, and + each handler self-scopes to the caller's teams and own keys + (_resolve_team_daily_activity_scope). self_managed_routes is the only list + granting them to a non-admin, and check_route_access is exact-match, so each + sub-path needs its own entry: dropping one 401s the dashboard before the + handler ever runs. + """ + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + def outcome() -> str: + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as exc: + return f"denied: {exc}" + return "allowed" + + assert outcome() == "allowed" + + @pytest.mark.parametrize( "user_role", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b066b3b80e6..38241926f8e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14646,6 +14646,218 @@ async def test_get_team_daily_activity_aggregated_rejects_bad_ranges( mock_aggregated.assert_not_called() +def _key_search_team_setup(mock_db_client, user_id: str, team_id: str): + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="test@example.com", + user_role="internal_user", + ) + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [Member(user_id=user_id, role="user")] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) + return mock_user_info + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_scopes_where_before_take(mock_db_client): + """A member's search must put the team and own-key scoping inside the same + Prisma where as the term, because `take` trims rows before Python sees them: + scoped outside the where, the top-N slice could be spent entirely on keys + the caller is not allowed to see.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + user_id = "test_user_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) + mock_user_info = _key_search_team_setup(mock_db_client, user_id, team_id) + + user_key_1 = MagicMock() + user_key_1.token = "user_key_1" + matched = MagicMock() + matched.token = "user_key_1" + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=[[user_key_1], [matched]]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=user_api_key_dict, + search="Needle", + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=480, + ) + + token_calls = mock_db_client.db.litellm_verificationtoken.find_many.call_args_list + assert len(token_calls) == 2 + search_kwargs = token_calls[1][1] + assert search_kwargs["where"] == { + "team_id": {"in": (team_id,)}, + "token": {"in": ("user_key_1",)}, + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ), + } + assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert search_kwargs["order"] == {"spend": "desc"} + + call_kwargs = mock_aggregated.call_args[1] + assert call_kwargs["api_key"] == ["user_key_1"] + assert call_kwargs["entity_id"] == [team_id] + assert call_kwargs["table_name"] == "litellm_dailyteamspend" + assert call_kwargs["include_entity_breakdown"] is True + assert call_kwargs["timezone_offset_minutes"] == 480 + assert call_kwargs["model"] is None + assert call_kwargs["entity_metadata_field"] == {team_id: {"team_alias": "Test Team"}} + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_admin_unscoped_where(mock_db_client): + """An admin's search has no caller scoping, so the where is the bare OR over + token, key alias and user id; every matched hash is passed through to the + aggregation.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + match_1 = MagicMock() + match_1.token = "h1" + match_2 = MagicMock() + match_2.token = "h2" + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[match_1, match_2]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=None, + ) + + search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + assert search_kwargs["where"] == { + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ) + } + assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert mock_aggregated.call_args[1]["api_key"] == ["h1", "h2"] + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_no_match_returns_empty_without_aggregating( + mock_db_client, +): + """A term matching no key still owes the caller the standard metadata shape + (api_key_limit, total_api_keys), and the aggregated query must not run.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + result = await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=None, + ) + + assert result.results == [] + assert result.metadata.total_api_keys == 0 + assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT + mock_aggregated.assert_not_called() + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_excludes_teams_in_where(mock_db_client): + """The dashboard always sends exclude_team_ids=litellm-dashboard; if that + filter stayed out of the where, matching keys in excluded teams could fill + the take=N slice and push visible matches out.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + matched = MagicMock() + matched.token = "h1" + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[matched]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids="litellm-dashboard", + timezone=None, + ) + + search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + assert search_kwargs["where"] == { + "team_id": {"notIn": ("litellm-dashboard",)}, + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ), + } + assert mock_aggregated.call_args[1]["exclude_entity_ids"] == ["litellm-dashboard"] + + def _wire_new_team_prisma(mock_db_client): mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b21108..16dc41c3ba8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -20,7 +20,7 @@ import type { ColumnDef } from "@tanstack/react-table"; import PaginationStatusAlerts from "@/components/shared/PaginationStatusAlerts"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import React, { type ReactNode, useMemo, useState } from "react"; +import React, { type ReactNode, useCallback, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; @@ -34,6 +34,7 @@ import { tagDailyActivityCall, teamDailyActivityAggregatedCall, teamDailyActivityCall, + teamDailyActivityKeySearchCall, userDailyActivityCall, } from "@/components/networking"; import { Logo } from "@/components/molecules/logo/Logo"; @@ -182,6 +183,16 @@ const EntityUsage: React.FC = ({ const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models"; const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []); + const searchTeamKeys = useCallback( + (query: string) => { + if (!accessToken || !startTime || !endTime) return Promise.resolve({}); + const teamIds = Array.isArray(entityFilterArg) ? entityFilterArg : null; + return teamDailyActivityKeySearchCall(accessToken, startTime, endTime, query, teamIds).then((data) => + processActivityData(data, "api_keys", teams || []), + ); + }, + [accessToken, startTime, endTime, entityFilterArg, teams], + ); const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {}; const getAllTags = () => { @@ -667,6 +678,7 @@ const EntityUsage: React.FC = ({ keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} apiKeyTruncation={apiKeyTruncation} + searchKeys={entityType === "team" ? searchTeamKeys : undefined} /> ), }, diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e6923a26d4c..edaf56a8a16 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1467,6 +1467,31 @@ export const teamDailyActivityAggregatedCall = async ( } }; +export const teamDailyActivityKeySearchCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + ...options: [search: string, teamIds?: string[] | null] +) => { + const [search, teamIds = null] = options; + try { + return await apiClient.get(`/team/daily/activity/aggregated/search`, { + accessToken, + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + search, + team_ids: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined, + exclude_team_ids: "litellm-dashboard", + }, + }); + } catch (error) { + console.error("Failed to search team daily activity keys:", error); + throw error; + } +}; + export type TeamUserSpendResponse = components["schemas"]["TeamUserSpendResponse"]; export const teamSpendByUserCall = async ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79d008c1b48..a398b2db93c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15645,6 +15645,27 @@ export interface paths { patch?: never; trace?: never; }; + "/team/daily/activity/aggregated/search": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Search Team Daily Activity Keys + * @description Aggregated daily team activity for the keys matching `search`, across every key the caller may + * see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend. + */ + get: operations["search_team_daily_activity_keys_team_daily_activity_aggregated_search_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/delete": { parameters: { query?: never; @@ -31378,7 +31399,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -66624,6 +66645,43 @@ export interface operations { }; }; }; + search_team_daily_activity_keys_team_daily_activity_aggregated_search_get: { + parameters: { + query: { + /** @description Exact token hash, or a case-insensitive substring of the key alias or owning user id */ + search: string; + team_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + exclude_team_ids?: string | null; + timezone?: number | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_team_team_delete_post: { parameters: { query?: never; From 1fbd1e9ce90580f801612d2016e9e2cc4ab553b5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:08:27 -0700 Subject: [PATCH 006/154] fix(bedrock): route unmapped openai family model ids to converse (#42713) * fix(bedrock): route unmapped openai family model ids to converse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): rename the e2e openai family backend constant global.openai.gpt-6-sol has a cost-map row now, so the constant no longer names an unmapped model --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/bedrock/common_utils.py | 2 ++ .../test_bedrock_provider_matrix_e2e.py | 19 ++++++++++++++++++- .../llms/bedrock/test_bedrock_common_utils.py | 19 ++++++++++++++++++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9b52f531cbb..5f044897b2c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1218,6 +1218,8 @@ class BedrockModelInfo(BaseLLMModelInfo): alt_model: Final = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) if base_model in litellm.bedrock_converse_models or alt_model in litellm.bedrock_converse_models: return "converse" + if _OPENAI_FAMILY_MODEL_RE.search(base_model): + return "converse" return "invoke" @staticmethod diff --git a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py index 5f0a931109c..21333d39849 100644 --- a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py @@ -8,7 +8,8 @@ caller can hand AWS support the request id behind a completion. Regional inference-profile ids are the deployment shape most Bedrock customers run; a v1.90.0 regression timed them out, and the Converse route keeps them covered in test_chat_completions_regression_e2e.py, so the invoke route carries its own -rows here. +rows here. The file also covers Bedrock-native OpenAI model ids taking the +default (Converse) route with max_tokens. """ from __future__ import annotations @@ -26,6 +27,7 @@ pytestmark = pytest.mark.e2e CONVERSE_REGIONAL_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" INVOKE_REGIONAL_BACKEND = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" +OPENAI_FAMILY_BACKEND = "bedrock/global.openai.gpt-6-sol" PROVIDER_HEADER_PREFIX = "llm_provider-" BEDROCK_REQUEST_ID_HEADER = "llm_provider-x-amzn-requestid" @@ -199,3 +201,18 @@ class TestBedrockInvokeRegionalModelIds: ) _assert_streamed_completion(result) + + +class TestBedrockOpenAIFamilyDefaultRoute: + @pytest.mark.covers("llm.chat_completions.bedrock_converse.basic.nonstream.works", exercised_on=[]) + def test_openai_family_model_id_completes_with_max_tokens( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model( + client, resources, "e2e-bedrock-openai-family", OPENAI_FAMILY_BACKEND + ) + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_prompt(), max_tokens=64))) + + _assert_completion(response) diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 117814a41ff..e5118f90e44 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -346,7 +346,7 @@ def test_route_prefix_matched_as_path_segment_not_substring(): BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" ) assert ( - BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "invoke" + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "converse" ) assert ( BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") @@ -964,3 +964,20 @@ def test_s3_static_key_pair_is_none_without_a_full_pair(partial_s3_pair): from litellm.llms.bedrock.common_utils import s3_static_key_pair assert s3_static_key_pair({"aws_access_key_id": "bedrock-key", **partial_s3_pair}) is None + + +def test_unmapped_openai_family_model_routes_to_converse(): + """A Bedrock-native OpenAI model that is not in the cost map yet must not fall to the invoke route. + + The invoke ``openai`` provider is the imported-model path and sends ``max_tokens``, which Bedrock + rejects for these models; Converse maps it to ``inferenceConfig.maxTokens``. + """ + from typing import Final + + import litellm + + unmapped: Final = "bedrock/global.openai.gpt-99-unmapped" + assert unmapped.removeprefix("bedrock/") not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "converse" + imported: Final = "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123" + assert BedrockModelInfo.get_bedrock_route(imported) == "openai" From e2781c47132a129b67667ed824d668899b4c3661 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:19:00 +0000 Subject: [PATCH 007/154] refactor(rust): move tests.rs files inline or under tests/ and drop autotests = false (#43028) * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): inline path-included test files into their owning src files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(rust): drop stray proptest regression file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust): cover lowercase, empty and non-authorization headers in bearer detection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/AGENTS.md | 10 + .../crates/callbacks-legacy-python/Cargo.toml | 1 - .../callbacks-legacy-python/src/adapter.rs | 1116 ++++- .../callbacks-legacy-python/src/deferred.rs | 150 +- .../crates/callbacks-legacy-python/src/lib.rs | 213 +- .../callbacks-legacy-python/tests/deferred.rs | 146 - .../tests/deployment_hooks.rs | 282 -- .../callbacks-legacy-python/tests/payload.rs | 523 --- .../callbacks-legacy-python/tests/support.rs | 205 - .../callbacks-legacy-python/tests/terminal.rs | 291 -- litellm-rust/crates/core/Cargo.toml | 1 - .../core/src/audio_transcription/mod.rs | 3 - .../crates/core/src/chat_completions/mod.rs | 3 - .../core/src/chat_completions/prepare.rs | 838 ++++ .../crates/core/src/chat_completions/tests.rs | 833 ---- .../crates/core/src/messages/common_utils.rs | 183 + litellm-rust/crates/core/src/messages/mod.rs | 3 - litellm-rust/crates/core/src/ocr/document.rs | 157 + litellm-rust/crates/core/src/ocr/mod.rs | 238 +- litellm-rust/crates/core/src/ocr/route.rs | 3586 +++++++++++++++++ .../tests.rs => tests/audio_transcription.rs} | 4 +- .../crates/core/tests/aws_textract_ocr.rs | 193 - .../crates/core/tests/azure_ai_ocr.rs | 293 -- .../tests/azure_document_intelligence_ocr.rs | 712 ---- litellm-rust/crates/core/tests/cohere_ocr.rs | 136 - .../crates/core/tests/deepseek_ocr.rs | 133 - .../messages/tests.rs => tests/messages.rs} | 132 +- litellm-rust/crates/core/tests/ocr.rs | 1033 ----- .../crates/core/tests/ocr/document.rs | 152 - litellm-rust/crates/core/tests/ocr/support.rs | 203 - litellm-rust/crates/core/tests/reducto_ocr.rs | 584 --- .../core/tests/vertex_ai_deepseek_ocr.rs | 143 - .../crates/core/tests/vertex_ai_ocr.rs | 293 -- litellm-rust/crates/http/src/request.rs | 12 + .../llms/src/anthropic/chat/transformation.rs | 4 - .../bedrock/chat/converse_transformation.rs | 4 - .../anthropic_chat_transformation.rs} | 18 +- .../bedrock_converse_transformation.rs} | 12 +- litellm-rust/crates/types/src/utils.rs | 2 +- 39 files changed, 6486 insertions(+), 6359 deletions(-) create mode 100644 litellm-rust/AGENTS.md delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/payload.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/support.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs delete mode 100644 litellm-rust/crates/core/src/chat_completions/tests.rs rename litellm-rust/crates/core/{src/audio_transcription/tests.rs => tests/audio_transcription.rs} (95%) delete mode 100644 litellm-rust/crates/core/tests/aws_textract_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_ai_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/cohere_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/deepseek_ocr.rs rename litellm-rust/crates/core/{src/messages/tests.rs => tests/messages.rs} (79%) delete mode 100644 litellm-rust/crates/core/tests/ocr.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/document.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/support.rs delete mode 100644 litellm-rust/crates/core/tests/reducto_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_ocr.rs rename litellm-rust/crates/llms/{src/anthropic/chat/tests.rs => tests/anthropic_chat_transformation.rs} (96%) rename litellm-rust/crates/llms/{src/bedrock/chat/tests.rs => tests/bedrock_converse_transformation.rs} (98%) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md new file mode 100644 index 00000000000..e5ffcd1c57a --- /dev/null +++ b/litellm-rust/AGENTS.md @@ -0,0 +1,10 @@ +# Rust workspace rules + +## Test placement + +- Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` +- A test that reaches private items lives inline, in a `#[cfg(test)] mod tests { ... }` at the bottom of the file that owns those items +- A test that only uses the crate's public API lives in `crates//tests/.rs`, next to `src/` +- Split a mixed test file along that line instead of widening visibility to move it +- A test for another crate's item belongs in that crate, not in a downstream one +- Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index fe19578e04d..ed5e0fb9691 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-host.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index e1742190205..75a635e9c63 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -487,11 +487,1115 @@ impl PythonLifecycle for LegacyLogging { } #[cfg(test)] -#[path = "../tests/deployment_hooks.rs"] -mod deployment_hooks_tests; +mod deployment_hooks_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::test_support::{legacy_call, local, namespace, run}; + + const CALL: &CStr = c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'logger': logger, 'document': document} +"; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn begin<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + asynchronous: bool, + ) -> (LegacyLogging, LifecycleStep) { + let mut logging = legacy_call(py, locals, asynchronous); + let kwargs = local(locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let step = logging.begin(py, kwargs, 0.0).unwrap(); + (logging, step) + } + + fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { + panic!("expected the prepared arguments"); + }; + arguments.into_bound(py) + } + + fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, CALL); + let (_, step) = begin(py, &locals, asynchronous); + assert_eq!(awaits_deployment_hook(&step), asynchronous); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); + }); + } + + #[test] + fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} +kwargs = {'logger': logger, 'document': document} +replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} +", + ); + let (mut logging, step) = begin(py, &locals, true); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) + .unwrap(); + locals.set_item("prepared", arguments(py, step)).unwrap(); + run( + py, + &locals, + c" +assert prepared['document'] is replacement +assert prepared['pages'] is replaced_kwargs['pages'] +assert prepared['litellm_logging_obj'] is logger +assert 'litellm_logging_obj' not in replaced_kwargs +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked is prepared +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); + } + + #[test] + fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +kwargs = {'logger': logger} +response = object() +replacement = object() +logger.hooks = {'pre': lambda kwargs: kwargs} +", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let step = logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replacement").unbind())) + .unwrap(); + let LifecycleStep::Response(returned) = step else { + panic!("expected the finalized response"); + }; + assert!(returned.bind(py).is(local(&locals, "replacement"))); + run( + py, + &locals, + c" +[finalized] = [value for name, value in logger.calls if name == 'finalize'] +assert finalized is replacement +", + ); + }); + } + + #[rstest] + #[case::pre_call(false)] + #[case::post_call(true)] + fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); + let (mut logging, _) = begin(py, &locals, true); + if post_call { + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + } + let cancellation = CancelledError::new_err("cancelled"); + let cancelled = cancellation.value(py).clone(); + let error = logging.resume(py, Err(cancellation)).err().unwrap(); + assert!(error.value(py).is(&cancelled)); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert!(!names.iter().any(|name| name.contains("handler"))); + }); + } + + #[rstest] + #[case::hook_completed(false)] + #[case::hook_cancelled(true)] + fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let failure = PyErr::from_value(local(&locals, "failure")); + let failed = LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Call, + error: &failure, + }; + let step = logging.emit(py, failed).unwrap(); + assert!(awaits_deployment_hook(&step)); + let hook_result = if cancelled { + Err(CancelledError::new_err("cancelled")) + } else { + Ok(py.None()) + }; + assert!(matches!( + logging.resume(py, hook_result).unwrap(), + LifecycleStep::Await(_) + )); + run( + py, + &locals, + c" +assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls +assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +class BudgetExceeded(Exception): + pass + +rejection = BudgetExceeded('over budget') + +class LimitedLogger(StubLogger): + def check_limits(self, arguments): + raise rejection + +logger = LimitedLogger() +logger.hooks = {'pre': lambda kwargs: kwargs} +kwargs = {'logger': logger} +", + ); + let mut logging = legacy_call(py, &locals, asynchronous); + let kwargs = local(&locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { + LifecycleStep::Await(_) => { + logging.resume(py, Ok(local(&locals, "kwargs").unbind())) + } + step => Ok(step), + }); + let error = result.err().unwrap(); + assert!(error.value(py).is(local(&locals, "rejection"))); + }); + } +} + #[cfg(test)] -#[path = "../tests/payload.rs"] -mod payload_tests; +mod payload_tests { + use std::ffi::CStr; + + use litellm_auth::SecretValue; + use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; + use proptest::prelude::*; + use pyo3::prelude::*; + use rstest::rstest; + use serde_json::{Map, Value, json}; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + /// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the + /// payload to the case's `on_pre_call`. + const PAYLOAD_LOGGER: &CStr = c" +class Request: + pass + +class PayloadLogger(StubLogger): + def update_from_kwargs(self, **update): + self.update = update + + def pre_call(self, input, api_key, additional_args): + self.record('pre_call', None) + self.pre = additional_args + self.pre_api_key = api_key + on_pre_call(additional_args) + + def post_call(self, original_response, api_key, additional_args): + self.record('post_call', None) + self.post = (original_response, api_key, additional_args) + +request = Request() +kwargs = {} +logger = PayloadLogger() +on_pre_call = lambda additional_args: None +check = lambda: None +"; + + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; + const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; + + fn document(source: &str) -> Value { + json!({"type": "document_url", "document_url": source}) + } + + fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) + } + + /// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with + /// the Python objects `script` binds, then delivers the provider's raw response the way the + /// driver does and runs the script's `check()`. + fn before_send_with_secrets( + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) + } + + /// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. + fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } + run(py, &locals, script); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let context = RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params, + secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), + }; + let wire = WireRequest { + url: "https://provider.invalid/ocr".into(), + headers: vec![("x-route".into(), "route".into())], + body, + }; + let step = logging.before_send(py, Box::new(wire), &context).unwrap(); + let raw = MachineEvent::ResponseReceived { + raw: RawResponse { + body: "raw response".into(), + }, + }; + assert!(matches!( + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done + )); + run(py, &locals, c"check()"); + let LifecycleStep::Wire(wire) = step else { + panic!("before_send did not hand back the wire request"); + }; + *wire + }) + } + + #[rstest] + #[case::caller_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +kwargs = {'document': document, 'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + #[case::request_attribute_behind_an_omitted_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +request.document = document +kwargs = {'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { + let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + } + + #[test] + fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' +def check(): + assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!(wire.body["document"], document(EDITED)); + } + + #[test] + fn a_body_key_the_route_rewrote_is_not_the_callers_object() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +kwargs = {'document': document} +observed = [] +def on_pre_call(args): + observed.append(args['complete_input_dict']['document'] is document) + args['complete_input_dict']['document']['document_name'] = 'edited.pdf' +def check(): + assert observed == [False], observed + assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body["document"], + json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) + ); + } + + #[test] + fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); + } + + #[rstest] + #[case::body( + c" +def on_pre_call(args): + args['complete_input_dict'] = {'replacement': True} +" + )] + #[case::headers( + c" +def on_pre_call(args): + args['headers'] = {'x-replacement': 'yes'} +" + )] + fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + + #[test] + fn pre_call_header_edit_reaches_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + args['headers']['x-callback'] = 'edited' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-callback".to_string(), "edited".to_string()), + ] + ); + } + + #[test] + fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { + let body = json!({"model": "model", "document": document(DOCUMENT)}); + before_send_with_secrets( + c" +logger_fn = lambda *args: None +kwargs = { + 'litellm_call_id': 'call-1', + 'client_secret': 'shh', + 'proxy_server_request': {'body': {}}, + 'logger_fn': logger_fn, + 'litellm_request_debug': True, + 'ocr_cost_per_page': 0.05, +} +observed = [] +on_pre_call = observed.append +def check(): + [args] = observed + assert args['api_base'] == 'https://provider.invalid/ocr', args + assert args['complete_input_dict'] == { + 'model': 'model', + 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, + }, args + update = logger.update + assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update + assert update['litellm_params']['litellm_call_id'] == 'call-1', update + assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update + assert update['litellm_params']['logger_fn'] is logger_fn, update + assert update['litellm_params']['litellm_request_debug'] is True, update + assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update + assert update['kwargs']['client_secret'] == '****', update + assert 'proxy_server_request' not in update['kwargs'], update + assert update['optional_params']['client_secret'] == '****', update +", + json!({"client_secret": "shh"}), + body, + &["client_secret"], + ); + } + + #[rstest] + #[case::added_key( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + #[case::replaced_document( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document'] = { + 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' + } +def check(): + assert document['document_url'] == 'data:application/pdf;base64,YWJj', document +", + json!({"document": document(EDITED)}) + )] + #[case::retained_body_edited_after_rebinding( + c" +def on_pre_call(args): + retained = args['complete_input_dict'] + args['complete_input_dict'] = {'rebound': True} + retained['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body); + assert_eq!(wire.body, expected); + } + + #[test] + fn retained_headers_edited_after_rebinding_reach_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + retained = args['headers'] + args['headers'] = {'x-rebound': 'rebound'} + retained['x-retained'] = 'sent' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-retained".to_string(), "sent".to_string()), + ] + ); + } + + #[test] + fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { + before_send( + c" +def check(): + original_response, api_key, additional_args = logger.post + assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args + assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] + assert additional_args['headers'] is logger.pre['headers'] +", + json!({"document": document(DOCUMENT)}), + ); + } + + #[test] + fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +def check(): + assert logger.names() == ['pre_call', 'post_call'], logger.calls +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body, + json!({"document": document(DOCUMENT), "include_image_base64": true}) + ); + } + + /// What one pre-call callback does to the payload it is handed. + #[derive(Clone, Debug)] + enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), + } + + impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } + } + + /// How the caller's keyword for a body key relates to what the route sends under it. + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, + } + + const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + + fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) + } + + fn key() -> impl Strategy { + "[a-z]{1,6}" + } + + fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] + } + + fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + } +} + #[cfg(test)] -#[path = "../tests/terminal.rs"] -mod terminal_tests; +mod terminal_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::PyRuntimeError; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { + LegacyLogging { + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), + ..legacy_call(py, locals, asynchronous) + } + } + + fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let response = local(locals, "response").unbind(); + logging + .emit( + py, + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, + ) + .unwrap() + } + + fn fail( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let failure = PyErr::from_value(local(locals, "failure")); + logging + .emit( + py, + LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Host, + error: &failure, + }, + ) + .unwrap() + } + + #[rstest] + #[case::sync_listened(false, c"", &["submit"])] + #[case::async_listened( + true, + c"", + &["async_success_handler", "enqueued", "sync_success_for_async_call"] + )] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn success_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c" +assert all(value is response for name, value in logger.calls if name.endswith('_handler')) +assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false, &["failure_handler"])] + #[case::asynchronous(true, &[])] + fn internal_calls_skip_failure_callbacks_only_when_asynchronous( + #[case] asynchronous: bool, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, asynchronous) + }; + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[test] + fn internal_async_calls_skip_the_async_success_fan_out() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, true) + }; + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", + ); + }); + } + + #[test] + fn a_failing_success_callback_is_reported_without_replacing_the_response() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +response = object() +failure = ValueError('terminal diagnostic') + +class FailingLogger(StubLogger): + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + assert!( + logging + .response + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "response")) + ); + run(py, &locals, c"assert unraisable_from(logger) == [failure]"); + }); + } + + #[rstest] + #[case::sync_listened(false, c"", &["failure_handler"])] + #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] + fn failure_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + let step = fail(py, &locals, &mut logging); + let awaits_async_handler = expected.contains(&"async_failure_handler"); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", + ); + }); + } + + #[test] + fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +failure = ValueError('selected') + +class FailingLogger(StubLogger): + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + raise RuntimeError('handler failed') + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Await(_) + )); + assert!( + logging + .error + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "failure")) + ); + run( + py, + &locals, + c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", + ); + }); + } + + #[rstest] + #[case::completed(None, true)] + #[case::handler_error(Some(false), true)] + #[case::cancelled(Some(true), false)] + fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( + #[case] error: Option, + #[case] done: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = logged(py, &locals, true); + fail(py, &locals, &mut logging); + let result = match error { + None => Ok(py.None()), + Some(false) => Err(PyRuntimeError::new_err("handler failed")), + Some(true) => Err(CancelledError::new_err("cancelled")), + }; + let expected = result.as_ref().err().map(|error| error.value(py).clone()); + match logging.resume(py, result) { + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), + Err(propagated) => { + assert!(!done); + assert!(propagated.value(py).is(expected.unwrap())); + } + } + }); + } + + #[test] + fn closing_restores_the_correlation_context_once() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c""); + let mut logging = logged(py, &locals, true); + logging.close(py); + logging.close(py); + run( + py, + &locals, + c"assert logger.names() == ['restore'], logger.calls", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index b18012f926e..292648ba69c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -63,5 +63,151 @@ impl PendingLogging { } #[cfg(test)] -#[path = "../tests/deferred.rs"] -mod tests; +mod tests { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::{PendingLogging, PendingSuccess}; + use crate::PythonLogger; + use crate::test_support::{local, namespace, run}; + + /// A deferred success for the namespace's `logger` and `response`, bound as `pending`. + fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: PythonLogger::new(local(&locals, "logger").unbind()), + response: Some(local(&locals, "response").unbind()), + start: py.None(), + end: Some(py.None()), + }), + }, + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + locals + } + + #[test] + fn release_enqueues_the_success_once_in_the_releasing_context() { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +from contextvars import ContextVar + +marker = ContextVar('marker', default='unset') +observed = [] + +def on_enqueue(coroutine): + observed.append(marker.get()) + pending.release(True) + +logger.on_enqueue = on_enqueue +", + ); + run( + py, + &locals, + c" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['release'], observed +assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.calls[0][1] is response +", + ); + }); + } + + #[test] + fn a_blocked_release_drops_the_success_for_good() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +pending.release(False) +pending.release(True) +assert logger.calls == [], logger.calls +", + ); + }); + } + + #[rstest] + #[case::ordinary_error(c"RuntimeError('queue full')", false)] + #[case::cancellation(c"asyncio.CancelledError()", true)] + fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( + #[case] failure: &CStr, + #[case] propagates: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +import asyncio + +def on_enqueue(coroutine): + raise failure + +logger.on_enqueue = on_enqueue +", + ); + locals + .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) + .unwrap(); + let released = local(&locals, "pending").call_method1("release", (true,)); + match released { + Ok(_) => assert!(!propagates), + Err(error) => { + assert!(propagates); + assert!(error.value(py).is(local(&locals, "failure"))); + } + } + locals.set_item("propagates", propagates).unwrap(); + run( + py, + &locals, + c" +pending.release(True) +assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +assert unraisable_from(logger) == ([] if propagates else [failure]) +", + ); + }); + } + + #[test] + fn an_unreleased_success_does_not_keep_its_logger_alive() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +import gc +import weakref + +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 44393792d1f..030bf03d4ba 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -16,13 +16,218 @@ mod deferred; mod logger; mod preparation; mod python; -#[cfg(test)] -#[path = "../tests/support.rs"] -mod test_support; - pub(crate) use adapter::LegacyLogging; pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; + +#[cfg(test)] +mod test_support { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyTuple}; + + use crate::{LegacyLogging, LegacySurface, PublicCall}; + + /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. + /// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python + /// signatures, and [`namespace`] binds every fake call against it. + pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + + /// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests + /// share one interpreter and run concurrently, so each fake is installed idempotently and + /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). + /// Every fake is bound against the contract first, so a call the real module would reject + /// fails here too. + const STUBS: &CStr = c" +import contextvars +import inspect +import json +import sys +import traceback +import types + +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): + sys.modules.setdefault(name, types.ModuleType(name)) + +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +CONTRACT = json.loads(python_contract) + + +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) + + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) + + +unraisable = sys.modules.setdefault( + 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') +) +if not hasattr(unraisable, 'events'): + unraisable.events = [] + sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) + + +def unraisable_from(owner): + return [error for source, error in unraisable.events if source is owner] + + +class StubCoroutine: + def __init__(self, logger): + self.logger = logger + + def enqueue(self): + self.logger.record('enqueued', None) + self.logger.on_enqueue(self) + + def close(self): + self.logger.record('closed', None) + + +class StubLogger: + def __init__(self): + self.calls = [] + self.hooks = {} + self.on_enqueue = lambda coroutine: None + + def record(self, name, value): + self.calls.append((name, value)) + + def names(self): + return [name for name, _ in self.calls] + + def hook(self, phase, value, call_type): + self.record(phase + '_hook', call_type) + return self.hooks.get(phase, lambda value: 'awaitable')(value) + + def check_limits(self, arguments): + self.record('check_limits', arguments) + + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + + def async_failure_handler(self, error, trace, start, end): + self.record('async_failure_handler', error) + return 'awaitable' + + def success_handler(self, response, start, end): + self.record('success_handler', response) + + def async_success_handler(self, response, start, end): + self.record('async_success_handler', response) + return StubCoroutine(self) + + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): + self.record('sync_success_for_async_call', response) + + +logger = StubLogger() +"; + + /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. + pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); + py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals + } + + pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { + py.run(code, Some(locals), Some(locals)).unwrap(); + } + + pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + locals.get_item(name).unwrap().unwrap() + } + + /// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). + pub(crate) fn legacy_call( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, + ) -> LegacyLogging { + let request = locals + .get_item("request") + .unwrap() + .unwrap_or_else(|| py.None().into_bound(py)); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .map(|kwargs| kwargs.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)); + let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + LegacyLogging::new( + py, + LegacySurface { + call_type: "test", + input_description: "test input", + stream: None, + }, + call, + asynchronous, + ) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs deleted file mode 100644 index 289ea1b2e7f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs +++ /dev/null @@ -1,146 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::{PendingLogging, PendingSuccess}; -use crate::PythonLogger; -use crate::test_support::{local, namespace, run}; - -/// A deferred success for the namespace's `logger` and `response`, bound as `pending`. -fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind()), - response: Some(local(&locals, "response").unbind()), - start: py.None(), - end: Some(py.None()), - }), - }, - ) - .unwrap(); - locals.set_item("pending", pending).unwrap(); - locals -} - -#[test] -fn release_enqueues_the_success_once_in_the_releasing_context() { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -from contextvars import ContextVar - -marker = ContextVar('marker', default='unset') -observed = [] - -def on_enqueue(coroutine): - observed.append(marker.get()) - pending.release(True) - -logger.on_enqueue = on_enqueue -", - ); - run( - py, - &locals, - c" -marker.set('release') -pending.release(True) -pending.release(True) -assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls -assert logger.calls[0][1] is response -", - ); - }); -} - -#[test] -fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls -", - ); - }); -} - -#[rstest] -#[case::ordinary_error(c"RuntimeError('queue full')", false)] -#[case::cancellation(c"asyncio.CancelledError()", true)] -fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( - #[case] failure: &CStr, - #[case] propagates: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -import asyncio - -def on_enqueue(coroutine): - raise failure - -logger.on_enqueue = on_enqueue -", - ); - locals - .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) - .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); - match released { - Ok(_) => assert!(!propagates), - Err(error) => { - assert!(propagates); - assert!(error.value(py).is(local(&locals, "failure"))); - } - } - locals.set_item("propagates", propagates).unwrap(); - run( - py, - &locals, - c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls -assert unraisable_from(logger) == ([] if propagates else [failure]) -", - ); - }); -} - -#[test] -fn an_unreleased_success_does_not_keep_its_logger_alive() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -import gc -import weakref - -logger.pending = pending -reference = weakref.ref(logger) -del logger, pending -gc.collect() -assert reference() is None -", - ); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs deleted file mode 100644 index 52c5e47f83f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::test_support::{legacy_call, local, namespace, run}; - -const CALL: &CStr = c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'logger': logger, 'document': document} -"; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn begin<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - asynchronous: bool, -) -> (LegacyLogging, LifecycleStep) { - let mut logging = legacy_call(py, locals, asynchronous); - let kwargs = local(locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let step = logging.begin(py, kwargs, 0.0).unwrap(); - (logging, step) -} - -fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { - let LifecycleStep::Arguments(arguments) = step else { - panic!("expected the prepared arguments"); - }; - arguments.into_bound(py) -} - -fn awaits_deployment_hook(step: &LifecycleStep) -> bool { - matches!(step, LifecycleStep::Await(_)) -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, CALL); - let (_, step) = begin(py, &locals, asynchronous); - assert_eq!(awaits_deployment_hook(&step), asynchronous); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); - }); -} - -#[test] -fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} -kwargs = {'logger': logger, 'document': document} -replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} -", - ); - let (mut logging, step) = begin(py, &locals, true); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) - .unwrap(); - locals.set_item("prepared", arguments(py, step)).unwrap(); - run( - py, - &locals, - c" -assert prepared['document'] is replacement -assert prepared['pages'] is replaced_kwargs['pages'] -assert prepared['litellm_logging_obj'] is logger -assert 'litellm_logging_obj' not in replaced_kwargs -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked is prepared -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( - #[case] asynchronous: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -opaque = object() -hooked = [] -logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} -kwargs = {'logger': logger, 'vendor_extension': opaque} -", - ); - let (mut logging, step) = begin(py, &locals, asynchronous); - let step = match step { - LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), - step => step, - }; - locals.set_item("prepared", arguments(py, step)).unwrap(); - locals.set_item("asynchronous", asynchronous).unwrap(); - run( - py, - &locals, - c" -assert prepared['vendor_extension'] is opaque -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked['vendor_extension'] is opaque -assert hooked == ([opaque] if asynchronous else []), hooked -", - ); - }); -} - -#[test] -fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -kwargs = {'logger': logger} -response = object() -replacement = object() -logger.hooks = {'pre': lambda kwargs: kwargs} -", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let step = logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replacement").unbind())) - .unwrap(); - let LifecycleStep::Response(returned) = step else { - panic!("expected the finalized response"); - }; - assert!(returned.bind(py).is(local(&locals, "replacement"))); - run( - py, - &locals, - c" -[finalized] = [value for name, value in logger.calls if name == 'finalize'] -assert finalized is replacement -", - ); - }); -} - -#[rstest] -#[case::pre_call(false)] -#[case::post_call(true)] -fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); - let (mut logging, _) = begin(py, &locals, true); - if post_call { - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - } - let cancellation = CancelledError::new_err("cancelled"); - let cancelled = cancellation.value(py).clone(); - let error = logging.resume(py, Err(cancellation)).err().unwrap(); - assert!(error.value(py).is(&cancelled)); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert!(!names.iter().any(|name| name.contains("handler"))); - }); -} - -#[rstest] -#[case::hook_completed(false)] -#[case::hook_cancelled(true)] -fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let failure = PyErr::from_value(local(&locals, "failure")); - let failed = LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &failure, - }; - let step = logging.emit(py, failed).unwrap(); - assert!(awaits_deployment_hook(&step)); - let hook_result = if cancelled { - Err(CancelledError::new_err("cancelled")) - } else { - Ok(py.None()) - }; - assert!(matches!( - logging.resume(py, hook_result).unwrap(), - LifecycleStep::Await(_) - )); - run( - py, - &locals, - c" -assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls -assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -class BudgetExceeded(Exception): - pass - -rejection = BudgetExceeded('over budget') - -class LimitedLogger(StubLogger): - def check_limits(self, arguments): - raise rejection - -logger = LimitedLogger() -logger.hooks = {'pre': lambda kwargs: kwargs} -kwargs = {'logger': logger} -", - ); - let mut logging = legacy_call(py, &locals, asynchronous); - let kwargs = local(&locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), - step => Ok(step), - }); - let error = result.err().unwrap(); - assert!(error.value(py).is(local(&locals, "rejection"))); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs b/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs deleted file mode 100644 index 5459b36af27..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs +++ /dev/null @@ -1,523 +0,0 @@ -use std::ffi::CStr; - -use litellm_auth::SecretValue; -use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; -use proptest::prelude::*; -use pyo3::prelude::*; -use rstest::rstest; -use serde_json::{Map, Value, json}; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the -/// payload to the case's `on_pre_call`. -const PAYLOAD_LOGGER: &CStr = c" -class Request: - pass - -class PayloadLogger(StubLogger): - def update_from_kwargs(self, **update): - self.update = update - - def pre_call(self, input, api_key, additional_args): - self.record('pre_call', None) - self.pre = additional_args - self.pre_api_key = api_key - on_pre_call(additional_args) - - def post_call(self, original_response, api_key, additional_args): - self.record('post_call', None) - self.post = (original_response, api_key, additional_args) - -request = Request() -kwargs = {} -logger = PayloadLogger() -on_pre_call = lambda additional_args: None -check = lambda: None -"; - -const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; -const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; - -fn document(source: &str) -> Value { - json!({"type": "document_url", "document_url": source}) -} - -fn before_send(script: &CStr, body: Value) -> WireRequest { - before_send_with_secrets(script, json!({}), body, &[]) -} - -/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with -/// the Python objects `script` binds, then delivers the provider's raw response the way the -/// driver does and runs the script's `check()`. -fn before_send_with_secrets( - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - before_send_bound(&[], script, optional_params, body, secret_fields) -} - -/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. -fn before_send_bound( - bindings: &[(&str, &Value)], - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, PAYLOAD_LOGGER); - for &(name, value) in bindings { - locals.set_item(name, to_py(py, value).unwrap()).unwrap(); - } - run(py, &locals, script); - let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), - ..legacy_call(py, &locals, false) - }; - let context = RequestContext { - model: "model".into(), - custom_llm_provider: "provider".into(), - optional_params, - secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), - api_key: Some(SecretValue::new("route-key")), - }; - let wire = WireRequest { - url: "https://provider.invalid/ocr".into(), - headers: vec![("x-route".into(), "route".into())], - body, - }; - let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = MachineEvent::ResponseReceived { - raw: RawResponse { - body: "raw response".into(), - }, - }; - assert!(matches!( - logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), - LifecycleStep::Done - )); - run(py, &locals, c"check()"); - let LifecycleStep::Wire(wire) = step else { - panic!("before_send did not hand back the wire request"); - }; - *wire - }) -} - -#[rstest] -#[case::caller_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -kwargs = {'document': document, 'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -#[case::request_attribute_behind_an_omitted_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -request.document = document -kwargs = {'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { - let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); -} - -#[test] -fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' -def check(): - assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!(wire.body["document"], document(EDITED)); -} - -#[test] -fn a_body_key_the_route_rewrote_is_not_the_callers_object() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -kwargs = {'document': document} -observed = [] -def on_pre_call(args): - observed.append(args['complete_input_dict']['document'] is document) - args['complete_input_dict']['document']['document_name'] = 'edited.pdf' -def check(): - assert observed == [False], observed - assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body["document"], - json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) - ); -} - -#[test] -fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { - let body = json!({"pages": [0]}); - let wire = before_send( - c" -opaque = object() -kwargs = {'pages': opaque} -observed = [] -on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) -def check(): - assert observed == [[0]], observed -", - body.clone(), - ); - assert_eq!(wire.body, body); -} - -#[rstest] -#[case::body( - c" -def on_pre_call(args): - args['complete_input_dict'] = {'replacement': True} -" -)] -#[case::headers( - c" -def on_pre_call(args): - args['headers'] = {'x-replacement': 'yes'} -" -)] -fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); - assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); -} - -#[test] -fn pre_call_header_edit_reaches_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - args['headers']['x-callback'] = 'edited' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-callback".to_string(), "edited".to_string()), - ] - ); -} - -#[test] -fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { - let body = json!({"model": "model", "document": document(DOCUMENT)}); - before_send_with_secrets( - c" -logger_fn = lambda *args: None -kwargs = { - 'litellm_call_id': 'call-1', - 'client_secret': 'shh', - 'proxy_server_request': {'body': {}}, - 'logger_fn': logger_fn, - 'litellm_request_debug': True, - 'ocr_cost_per_page': 0.05, -} -observed = [] -on_pre_call = observed.append -def check(): - [args] = observed - assert args['api_base'] == 'https://provider.invalid/ocr', args - assert args['complete_input_dict'] == { - 'model': 'model', - 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, - }, args - update = logger.update - assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update - assert update['litellm_params']['litellm_call_id'] == 'call-1', update - assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update - assert update['litellm_params']['logger_fn'] is logger_fn, update - assert update['litellm_params']['litellm_request_debug'] is True, update - assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update - assert update['kwargs']['client_secret'] == '****', update - assert 'proxy_server_request' not in update['kwargs'], update - assert update['optional_params']['client_secret'] == '****', update -", - json!({"client_secret": "shh"}), - body, - &["client_secret"], - ); -} - -#[rstest] -#[case::added_key( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -#[case::replaced_document( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document'] = { - 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' - } -def check(): - assert document['document_url'] == 'data:application/pdf;base64,YWJj', document -", - json!({"document": document(EDITED)}) -)] -#[case::retained_body_edited_after_rebinding( - c" -def on_pre_call(args): - retained = args['complete_input_dict'] - args['complete_input_dict'] = {'rebound': True} - retained['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body); - assert_eq!(wire.body, expected); -} - -#[test] -fn retained_headers_edited_after_rebinding_reach_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - retained = args['headers'] - args['headers'] = {'x-rebound': 'rebound'} - retained['x-retained'] = 'sent' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-retained".to_string(), "sent".to_string()), - ] - ); -} - -#[test] -fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { - before_send( - c" -def check(): - original_response, api_key, additional_args = logger.post - assert original_response == 'raw response', original_response - assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) - assert additional_args == { - 'complete_input_dict': logger.pre['complete_input_dict'], - 'headers': logger.pre['headers'], - }, additional_args - assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] -", - json!({"document": document(DOCUMENT)}), - ); -} - -#[test] -fn every_request_runs_the_full_pre_call_and_post_call() { - let wire = before_send( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -def check(): - assert logger.names() == ['pre_call', 'post_call'], logger.calls -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body, - json!({"document": document(DOCUMENT), "include_image_base64": true}) - ); -} - -/// What one pre-call callback does to the payload it is handed. -#[derive(Clone, Debug)] -enum Edit { - Nothing, - Set(String, Value), - Remove(String), - Rebind(Value), - RebindThenSetRetained(String, Value), -} - -impl Edit { - fn script(&self) -> Value { - match self { - Self::Nothing => json!({"kind": "nothing"}), - Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), - Self::Remove(key) => json!({"kind": "remove", "key": key}), - Self::Rebind(value) => json!({"kind": "rebind", "value": value}), - Self::RebindThenSetRetained(key, value) => { - json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) - } - } - } - - /// The legacy contract: the provider is sent the body object `pre_call` received, as - /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and - /// leaves that object alone. - fn sent(&self, body: &Map) -> Value { - let mut sent = body.clone(); - match self { - Self::Nothing | Self::Rebind(_) => {} - Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { - sent.insert(key.clone(), value.clone()); - } - Self::Remove(key) => { - sent.remove(key); - } - } - Value::Object(sent) - } -} - -/// How the caller's keyword for a body key relates to what the route sends under it. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Caller { - PassedUnchanged, - RewrittenByTheRoute, - NotPassed, -} - -const MODEL: &CStr = c" -aliased = {} -def on_pre_call(args): - body = args['complete_input_dict'] - aliased.update({name: body[name] is kwargs[name] for name in unchanged}) - kind = edit['kind'] - if kind == 'set': - body[edit['key']] = edit['value'] - elif kind == 'remove': - body.pop(edit['key'], None) - elif kind == 'rebind': - args['complete_input_dict'] = edit['value'] - elif kind == 'rebind_then_set_retained': - args['complete_input_dict'] = {} - body[edit['key']] = edit['value'] -def check(): - assert aliased == {name: True for name in unchanged}, aliased - assert logger.names() == ['pre_call', 'post_call'], logger.calls -"; - -fn json_value() -> impl Strategy { - let leaf = prop_oneof![ - Just(Value::Null), - any::().prop_map(Value::from), - any::().prop_map(Value::from), - any::() - .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) - .prop_map(Value::from), - ".{0,8}".prop_map(Value::from), - ]; - leaf.prop_recursive(3, 24, 4, |inner| { - prop_oneof![ - prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), - prop::collection::btree_map(key(), inner, 0..4) - .prop_map(|fields| Value::Object(fields.into_iter().collect())), - ] - }) -} - -fn key() -> impl Strategy { - "[a-z]{1,6}" -} - -fn caller() -> impl Strategy { - prop_oneof![ - Just(Caller::PassedUnchanged), - Just(Caller::RewrittenByTheRoute), - Just(Caller::NotPassed), - ] -} - -fn edit() -> impl Strategy { - prop_oneof![ - Just(Edit::Nothing), - (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), - key().prop_map(Edit::Remove), - json_value().prop_map(Edit::Rebind), - (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), - ] -} - -proptest! { - #![proptest_config(ProptestConfig::with_cases(128))] - - /// For any body, any caller keywords and any callback edit: every keyword the route - /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is - /// sent exactly what the model says, so a callback that edits nothing changes nothing. - #[test] - fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( - fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), - edit in edit(), - ) { - let body: Map = fields - .iter() - .map(|(name, (value, _))| (name.clone(), value.clone())) - .collect(); - let kwargs: Map = fields - .iter() - .filter_map(|(name, (value, caller))| match caller { - Caller::PassedUnchanged => Some((name.clone(), value.clone())), - Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), - Caller::NotPassed => None, - }) - .collect(); - let unchanged: Value = fields - .iter() - .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) - .map(|(name, _)| Value::from(name.clone())) - .collect(); - - let wire = before_send_bound( - &[ - ("kwargs", &Value::Object(kwargs)), - ("unchanged", &unchanged), - ("edit", &edit.script()), - ], - MODEL, - json!({}), - Value::Object(body.clone()), - &[], - ); - - prop_assert_eq!(wire.body, edit.sent(&body)); - prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); - } -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs b/litellm-rust/crates/callbacks-legacy-python/tests/support.rs deleted file mode 100644 index d0c02fa6da5..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; - -use crate::{LegacyLogging, LegacySurface, PublicCall}; - -/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. -/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python -/// signatures, and [`namespace`] binds every fake call against it. -pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); - -/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests -/// share one interpreter and run concurrently, so each fake is installed idempotently and -/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). -/// Every fake is bound against the contract first, so a call the real module would reject -/// fails here too. -const STUBS: &CStr = c" -import contextvars -import inspect -import json -import sys -import traceback -import types - -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): - sys.modules.setdefault(name, types.ModuleType(name)) - -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] -CONTRACT = json.loads(python_contract) - - -def contracted(name, fake): - signature = inspect.Signature( - [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] - ) - - def checked(*args, **kwargs): - signature.bind(*args, **kwargs) - return fake(*args, **kwargs) - - return checked - - -if not hasattr(legacy, 'is_internal'): - legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) - -FAKES = { - 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - ), - 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), - 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=provider, - ), - 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( - original_response, api_key, additional_args - ), - 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), - 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( - response, start, end - ), - 'failure_handler': lambda logger, error, start, end, asynchronous: ( - logger.async_failure_handler if asynchronous else logger.failure_handler - )(error, ''.join(traceback.format_exception(error)), start, end), - 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), - 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), - 'enqueue_logging': lambda coroutine: coroutine.enqueue(), - 'restore_context': lambda logger: logger.record('restore', None), - 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), - 'is_internal_call': lambda: legacy.is_internal.get(), - 'credential_list': lambda: [], - 'warn_unknown_credential': lambda name, loaded: None, - 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), - 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( - 'success', response, call_type - ), - 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger: logger.record('stream_opened', None), - 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( - 'stream_success', list(chunks) - ), - 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), -} -assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) -for name, fake in FAKES.items(): - setattr(legacy, name, contracted(name, fake)) - - -unraisable = sys.modules.setdefault( - 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') -) -if not hasattr(unraisable, 'events'): - unraisable.events = [] - sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) - - -def unraisable_from(owner): - return [error for source, error in unraisable.events if source is owner] - - -class StubCoroutine: - def __init__(self, logger): - self.logger = logger - - def enqueue(self): - self.logger.record('enqueued', None) - self.logger.on_enqueue(self) - - def close(self): - self.logger.record('closed', None) - - -class StubLogger: - def __init__(self): - self.calls = [] - self.hooks = {} - self.on_enqueue = lambda coroutine: None - - def record(self, name, value): - self.calls.append((name, value)) - - def names(self): - return [name for name, _ in self.calls] - - def hook(self, phase, value, call_type): - self.record(phase + '_hook', call_type) - return self.hooks.get(phase, lambda value: 'awaitable')(value) - - def check_limits(self, arguments): - self.record('check_limits', arguments) - - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - - def async_failure_handler(self, error, trace, start, end): - self.record('async_failure_handler', error) - return 'awaitable' - - def success_handler(self, response, start, end): - self.record('success_handler', response) - - def async_success_handler(self, response, start, end): - self.record('async_success_handler', response) - return StubCoroutine(self) - - def handle_sync_success_callbacks_for_async_calls(self, response, start, end): - self.record('sync_success_for_async_call', response) - - -logger = StubLogger() -"; - -/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. -pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); - py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); - py.run(script, Some(&locals), Some(&locals)).unwrap(); - locals -} - -pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { - py.run(code, Some(locals), Some(locals)).unwrap(); -} - -pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - locals.get_item(name).unwrap().unwrap() -} - -/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). -pub(crate) fn legacy_call( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - asynchronous: bool, -) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); - LegacyLogging::new( - py, - LegacySurface { - call_type: "test", - input_description: "test input", - stream: None, - }, - call, - asynchronous, - ) -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs deleted file mode 100644 index f68209233f2..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs +++ /dev/null @@ -1,291 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { - LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind())), - ..legacy_call(py, locals, asynchronous) - } -} - -fn succeed( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - logging: &mut LegacyLogging, -) -> LifecycleStep { - let response = local(locals, "response").unbind(); - logging - .emit( - py, - LifecycleEvent::Succeeded { - timing: TIMING, - response: &response, - }, - ) - .unwrap() -} - -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { - let failure = PyErr::from_value(local(locals, "failure")); - logging - .emit( - py, - LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Host, - error: &failure, - }, - ) - .unwrap() -} - -#[rstest] -#[case::sync_listened(false, c"", &["submit"])] -#[case::async_listened( - true, - c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] -)] -#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] -#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c" -assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false, &["failure_handler"])] -#[case::asynchronous(true, &[])] -fn internal_calls_skip_failure_callbacks_only_when_asynchronous( - #[case] asynchronous: bool, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, asynchronous) - }; - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - }); -} - -#[test] -fn internal_async_calls_skip_the_async_success_fan_out() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, true) - }; - succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); - }); -} - -#[test] -fn a_failing_success_callback_is_reported_without_replacing_the_response() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -response = object() -failure = ValueError('terminal diagnostic') - -class FailingLogger(StubLogger): - def handle_sync_success_callbacks_for_async_calls(self, *args): - raise failure - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - assert!( - logging - .response - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "response")) - ); - run(py, &locals, c"assert unraisable_from(logger) == [failure]"); - }); -} - -#[rstest] -#[case::sync_listened(false, c"", &["failure_handler"])] -#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -fn failure_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - let step = fail(py, &locals, &mut logging); - let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!( - matches!(step, LifecycleStep::Await(_)), - awaits_async_handler - ); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", - ); - }); -} - -#[test] -fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -failure = ValueError('selected') - -class FailingLogger(StubLogger): - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - raise RuntimeError('handler failed') - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Await(_) - )); - assert!( - logging - .error - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "failure")) - ); - run( - py, - &locals, - c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", - ); - }); -} - -#[rstest] -#[case::completed(None, true)] -#[case::handler_error(Some(false), true)] -#[case::cancelled(Some(true), false)] -fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( - #[case] error: Option, - #[case] done: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = logged(py, &locals, true); - fail(py, &locals, &mut logging); - let result = match error { - None => Ok(py.None()), - Some(false) => Err(PyRuntimeError::new_err("handler failed")), - Some(true) => Err(CancelledError::new_err("cancelled")), - }; - let expected = result.as_ref().err().map(|error| error.value(py).clone()); - match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), - Err(propagated) => { - assert!(!done); - assert!(propagated.value(py).is(expected.unwrap())); - } - } - }); -} - -#[test] -fn closing_restores_the_correlation_context_once() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c""); - let mut logging = logged(py, &locals, true); - logging.close(py); - logging.close(py); - run( - py, - &locals, - c"assert logger.names() == ['restore'], logger.calls", - ); - }); -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 4626781f3e3..d7096cdd774 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-secrets.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index af9c398c065..801fd5e9673 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -14,6 +14,3 @@ pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Resu execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 81d35044d08..224c9d8cfed 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -51,6 +51,3 @@ pub fn chat_completions_decline_reason( .unsupported_reason(&messages, optional_params) .map(|reason| reason.0) } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index c8e6365121e..afea46221f5 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -143,3 +143,841 @@ pub(super) fn prepare_provider_request( timeout: request.timeout, }) } + +#[cfg(test)] +mod tests { + use litellm_llms::base_llm::chat::transformation::RequestAuth; + use serde_json::{Map, Value, json}; + + use super::{prepare_provider_request, resolve_request}; + use crate::chat_completions::{ + Error, + types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}, + }; + + fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, + ) -> Result { + prepare_provider_request(resolve_request(request)?) + } + + fn request<'a>( + model: &'a str, + provider: Option<&'a str>, + messages: Value, + optional_params: Value, + ) -> ChatCompletionsRequest<'a> { + ChatCompletionsRequest { + model, + messages, + optional_params: match optional_params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: provider, + extra_headers: None, + timeout: None, + } + } + + /// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers + /// carry resolved credentials), so unwrap the failure case by hand. + fn decline(request: ChatCompletionsRequest<'_>) -> Error { + match prepare_chat_completions_call(request) { + Err(error) => error, + Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), + } + } + + #[test] + fn resolves_the_provider_from_the_model_prefix() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); + assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); + } + + #[test] + fn strips_an_explicit_provider_prefix_from_the_model() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + } + + #[test] + fn adds_the_auth_and_default_headers() { + let prepared = prepare_chat_completions_call(request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert!( + prepared + .upstream_headers + .contains(&("x-api-key".to_string(), "sk-test".to_string())) + ); + assert!( + prepared + .upstream_headers + .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) + ); + assert!(matches!( + prepared.auth, + RequestAuth::Header { + name: "x-api-key", + .. + } + )); + } + + #[test] + fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { + // Python builds `{**headers, **anthropic_headers}`, so the deployment's key + // overwrites a forwarded one. Honouring the caller's would let whoever sends + // the request choose the Anthropic principal it bills to. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "X-Api-Key".to_string(), + json!("sk-caller"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + } + + #[test] + fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { + // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` + // for an OAuth token, so re-adding the key here would put the credential into + // a header the host removed on purpose. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ( + "Authorization".to_string(), + json!("Bearer sk-ant-oat01-token"), + ), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), + "the resolved key must not be applied over an OAuth bearer, got {:?}", + prepared.upstream_headers + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-token") + ); + } + + #[test] + fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { + // Only an OAuth bearer replaces the credential. Python sends the deployment's + // `x-api-key` alongside any other forwarded `authorization`, so deferring on + // the mere presence of that header would drop the deployment's auth. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ("Authorization".to_string(), json!("Bearer unrelated")), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer unrelated"), + "the unrelated authorization must survive, got {:?}", + prepared.upstream_headers + ); + } + + #[test] + fn declines_an_unsupported_request_before_resolving_credentials() { + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ); + call.api_key = None; + // No api_key is set and no env is consulted: the gate must run first, so the + // error is the decline rather than a missing-credential error. + assert_eq!(decline(call), Error::Unsupported("streaming")); + } + + #[test] + fn rejects_an_unknown_provider() { + assert_eq!( + decline(request( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider("openai".to_string()) + ); + } + + #[test] + fn rejects_a_model_with_no_resolvable_provider() { + assert!(matches!( + decline(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider(_) + )); + } + + #[test] + fn rejects_an_empty_or_malformed_message_list() { + assert_eq!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!([]), + json!({}), + )), + Error::InvalidRequest("chat completions requires at least one message".to_string()) + ); + assert!(matches!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!("not a list"), + json!({}), + )), + Error::InvalidRequest(_) + )); + } + + #[test] + fn rejects_non_string_extra_headers() { + let mut call = request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); + assert_eq!( + decline(call), + Error::Headers(litellm_http::request::HeaderError { + context: "chat completions", + name: "x-trace".to_string(), + actual: "number", + }) + ); + } + + #[test] + fn prepares_a_bedrock_call_without_resolving_credentials() { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.api_key = None; + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert_eq!( + prepared.url, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); + assert_eq!( + prepared.auth, + RequestAuth::AwsSigV4 { + region: "us-east-1".to_string(), + service: "bedrock", + } + ); + // SigV4 signs the serialized body, so prepare must not have added an + // Authorization header; the handler does it. + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) + ); + assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); + } + + #[tokio::test] + async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { + // Python signs only the AWS header set and reattaches the rest, so a header + // the caller forwarded rides along without joining the canonical request. + // Signing it makes Converse 403 on a deployment that works on Python. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + // A key would resolve to a bearer token and never reach the signer. + call.api_key = None; + call.extra_headers = Some(Map::from_iter([( + "x-request-id".to_string(), + json!("abc-123"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let signed = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect("signs"); + + let authorization = signed + .header("authorization") + .expect("carries an authorization header") + .to_string(); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256"), + "expected a SigV4 signature, got {authorization}" + ); + assert!( + !authorization.contains("x-request-id"), + "forwarded header reached SignedHeaders: {authorization}" + ); + // It still goes on the wire, it is just not part of the signature. + assert!( + signed + .headers() + .iter() + .any(|(name, value)| name == "x-request-id" && value == "abc-123"), + "forwarded header was dropped instead of reattached" + ); + } + + #[tokio::test] + async fn a_forwarded_header_the_signer_computes_declines_to_python() { + // Reattaching the caller's copy next to the computed one puts the name on + // the wire twice and Bedrock rejects the pair, so a request carrying one + // has to go to Python instead of being signed here. + for forwarded in [ + "Authorization", + "x-amz-date", + "x-amz-security-token", + "Date", + ] { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + call.api_key = None; + call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let error = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect_err("{forwarded} should decline instead of being signed"); + assert!( + matches!(error, Error::Unsupported(_)), + "{forwarded} declined as {error:?}, which the host would not fall back on" + ); + } + } + + #[test] + fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { + // `get_request_headers` assigns `headers["Authorization"]` unconditionally + // once a bearer token resolves, so the deployment's identity wins on + // Python. Keeping the caller's would authorize and bill the call as a + // different principal, and only when the deployment carries `rust: true`. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.extra_headers = Some(Map::from_iter([( + "Authorization".to_string(), + json!("Bearer caller-supplied"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let authorizations: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + authorizations, + vec!["Bearer sk-test"], + "the deployment token must be the only authorization on the wire" + ); + } + + #[test] + fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { + // The opposite precedence, and deliberate: Anthropic's own transform + // honours a forwarded OAuth bearer, so the Bedrock fix above must not be + // generalized into a rule that the configured key always wins. + // + // An OAuth bearer is the whole of that exception. This forwarded a plain + // `x-api-key` until round 17, which read as the same claim and was not: + // Python overwrites a forwarded `x-api-key` with the deployment's. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "authorization".to_string(), + json!("Bearer sk-ant-oat01-forwarded"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .map(|(_, value)| value.as_str()) + .collect(); + assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-forwarded") + ); + } + + #[test] + fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { + // The configured bearer identity has its own account and quota boundary, + // so a request carrying one must not be signed as whatever principal the + // host's AWS credentials resolve to. + let prepared = prepare_chat_completions_call(request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + )) + .expect("prepares"); + assert_eq!( + prepared.auth, + RequestAuth::Bearer { + token: "sk-test".to_string() + } + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-test"), + "prepare did not carry the bearer token" + ); + } + + fn decline_reason( + model: &str, + provider: Option<&str>, + messages: Value, + params: Value, + ) -> Option<&'static str> { + let params = match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }; + crate::chat_completions::chat_completions_decline_reason(model, provider, messages, ¶ms) + } + + #[test] + fn the_gate_accepts_what_prepare_accepts() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ), + None + ); + } + + #[test] + fn the_gate_declines_without_resolving_credentials_or_calling_out() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ), + Some("streaming") + ); + assert_eq!( + decline_reason( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!("nope"), + json!({}) + ), + Some("unreadable message list") + ); + assert_eq!( + decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), + Some("empty message list") + ); + } + + #[test] + fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { + // A gate that accepts what prepare then declines would make the host emit + // its pre-call logging on a path that falls back, so pin the agreement. + for (messages, params) in [ + ( + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 8}), + ), + ( + json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), + json!({"temperature": 0.1}), + ), + ( + json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), + json!({}), + ), + ] { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params.clone() + ), + None, + "gate declined {messages}" + ); + prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params, + )) + .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); + } + } + + mod round_trip { + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::*; + use crate::chat_completions::chat_completions; + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") + { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + fn http_response(status: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + } + + /// Serve one request from a stub upstream and hand back what it received. + async fn serve_once( + status: &'static str, + body: &'static str, + ) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let port = listener.local_addr().expect("addr").port(); + let handle = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts"); + let received = read_http_request(&mut socket).await; + socket + .write_all(http_response(status, body).as_bytes()) + .await + .expect("writes response"); + socket.flush().await.expect("flushes"); + received + }); + (format!("http://127.0.0.1:{port}/v1/messages"), handle) + } + + fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages, + optional_params: match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: Some(std::time::Duration::from_secs(10)), + } + } + + const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + #[tokio::test] + async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { + let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; + let response = chat_completions(call( + &api_base, + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"max_tokens": 16}), + )) + .await + .expect("call succeeds"); + + let received = handle.await.expect("server task"); + let sent: Value = serde_json::from_str( + received + .split_once("\r\n\r\n") + .expect("request has a body") + .1, + ) + .expect("body is json"); + assert_eq!( + sent["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + sent["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(sent["max_tokens"], json!(16)); + assert!(received.to_lowercase().contains("x-api-key: sk-test")); + + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { + // The provider was called and billed, so the host must not retry this + // on its own path. `MissingField` here would read as a pre-send + // decline and be retried; `InvalidResponse` cannot. + const NO_USAGE: &str = + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; + let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { + const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; + let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn an_upstream_error_status_keeps_its_code() { + let (api_base, handle) = + serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("upstream rejects"); + handle.await.expect("server task"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) + ), + "expected a 429, got {err:?}" + ); + } + + #[tokio::test] + async fn a_connection_that_is_never_established_declines_instead_of_failing() { + // Nothing was sent, so nothing was billed and the host can still serve + // the request. Classing this with the post-send failures would turn a + // recoverable fallback into a user-facing error on exactly the + // deployments whose transport is configured only on the Python client. + let port = { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + listener.local_addr().expect("has an address").port() + // Dropped here, so the port is closed and the connect is refused. + }; + let err = chat_completions(call( + &format!("http://127.0.0.1:{port}/v1/messages"), + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("nothing is listening"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Connect(_)) + ), + "expected a pre-send connect failure, got {err:?}" + ); + } + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + use crate::chat_completions::handler::as_response_error; + + for original in [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth(litellm_auth::Error::InvalidHeader), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), Error::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + // An upstream status is already unambiguous, so it survives intact. + assert!(matches!( + as_response_error(Error::Transport(litellm_http::transport::Error::Http { + status: 500, + body: "boom".to_string() + })), + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) + )); + } + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs deleted file mode 100644 index dd5938cf168..00000000000 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ /dev/null @@ -1,833 +0,0 @@ -use litellm_llms::base_llm::chat::transformation::RequestAuth; -use serde_json::{Map, Value, json}; - -use super::{ - Error, - prepare::{prepare_provider_request, resolve_request}, -}; -use crate::chat_completions::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; - -fn prepare_chat_completions_call( - request: ChatCompletionsRequest<'_>, -) -> Result { - prepare_provider_request(resolve_request(request)?) -} - -fn request<'a>( - model: &'a str, - provider: Option<&'a str>, - messages: Value, - optional_params: Value, -) -> ChatCompletionsRequest<'a> { - ChatCompletionsRequest { - model, - messages, - optional_params: match optional_params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: None, - custom_llm_provider: provider, - extra_headers: None, - timeout: None, - } -} - -/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers -/// carry resolved credentials), so unwrap the failure case by hand. -fn decline(request: ChatCompletionsRequest<'_>) -> Error { - match prepare_chat_completions_call(request) { - Err(error) => error, - Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), - } -} - -#[test] -fn resolves_the_provider_from_the_model_prefix() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); - assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); - assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); -} - -#[test] -fn strips_an_explicit_provider_prefix_from_the_model() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); -} - -#[test] -fn adds_the_auth_and_default_headers() { - let prepared = prepare_chat_completions_call(request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert!( - prepared - .upstream_headers - .contains(&("x-api-key".to_string(), "sk-test".to_string())) - ); - assert!( - prepared - .upstream_headers - .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) - ); - assert!(matches!( - prepared.auth, - RequestAuth::Header { - name: "x-api-key", - .. - } - )); -} - -#[test] -fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { - // Python builds `{**headers, **anthropic_headers}`, so the deployment's key - // overwrites a forwarded one. Honouring the caller's would let whoever sends - // the request choose the Anthropic principal it bills to. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "X-Api-Key".to_string(), - json!("sk-caller"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); -} - -#[test] -fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { - // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` - // for an OAuth token, so re-adding the key here would put the credential into - // a header the host removed on purpose. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ( - "Authorization".to_string(), - json!("Bearer sk-ant-oat01-token"), - ), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), - "the resolved key must not be applied over an OAuth bearer, got {:?}", - prepared.upstream_headers - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-token") - ); -} - -#[test] -fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { - // Only an OAuth bearer replaces the credential. Python sends the deployment's - // `x-api-key` alongside any other forwarded `authorization`, so deferring on - // the mere presence of that header would drop the deployment's auth. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ("Authorization".to_string(), json!("Bearer unrelated")), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer unrelated"), - "the unrelated authorization must survive, got {:?}", - prepared.upstream_headers - ); -} - -#[test] -fn declines_an_unsupported_request_before_resolving_credentials() { - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ); - call.api_key = None; - // No api_key is set and no env is consulted: the gate must run first, so the - // error is the decline rather than a missing-credential error. - assert_eq!(decline(call), Error::Unsupported("streaming")); -} - -#[test] -fn rejects_an_unknown_provider() { - assert_eq!( - decline(request( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider("openai".to_string()) - ); -} - -#[test] -fn rejects_a_model_with_no_resolvable_provider() { - assert!(matches!( - decline(request( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider(_) - )); -} - -#[test] -fn rejects_an_empty_or_malformed_message_list() { - assert_eq!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!([]), - json!({}), - )), - Error::InvalidRequest("chat completions requires at least one message".to_string()) - ); - assert!(matches!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!("not a list"), - json!({}), - )), - Error::InvalidRequest(_) - )); -} - -#[test] -fn rejects_non_string_extra_headers() { - let mut call = request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); - assert_eq!( - decline(call), - Error::Headers(litellm_http::request::HeaderError { - context: "chat completions", - name: "x-trace".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn prepares_a_bedrock_call_without_resolving_credentials() { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.api_key = None; - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert_eq!( - prepared.url, - "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" - ); - assert_eq!( - prepared.auth, - RequestAuth::AwsSigV4 { - region: "us-east-1".to_string(), - service: "bedrock", - } - ); - // SigV4 signs the serialized body, so prepare must not have added an - // Authorization header; the handler does it. - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) - ); - assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); -} - -#[tokio::test] -async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { - // Python signs only the AWS header set and reattaches the rest, so a header - // the caller forwarded rides along without joining the canonical request. - // Signing it makes Converse 403 on a deployment that works on Python. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - // A key would resolve to a bearer token and never reach the signer. - call.api_key = None; - call.extra_headers = Some(Map::from_iter([( - "x-request-id".to_string(), - json!("abc-123"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = super::handler::outbound_request(&prepared) - .await - .expect("signs"); - - let authorization = signed - .header("authorization") - .expect("carries an authorization header") - .to_string(); - assert!( - authorization.starts_with("AWS4-HMAC-SHA256"), - "expected a SigV4 signature, got {authorization}" - ); - assert!( - !authorization.contains("x-request-id"), - "forwarded header reached SignedHeaders: {authorization}" - ); - // It still goes on the wire, it is just not part of the signature. - assert!( - signed - .headers() - .iter() - .any(|(name, value)| name == "x-request-id" && value == "abc-123"), - "forwarded header was dropped instead of reattached" - ); -} - -#[tokio::test] -async fn a_forwarded_header_the_signer_computes_declines_to_python() { - // Reattaching the caller's copy next to the computed one puts the name on - // the wire twice and Bedrock rejects the pair, so a request carrying one - // has to go to Python instead of being signed here. - for forwarded in [ - "Authorization", - "x-amz-date", - "x-amz-security-token", - "Date", - ] { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - call.api_key = None; - call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = super::handler::outbound_request(&prepared) - .await - .expect_err("{forwarded} should decline instead of being signed"); - assert!( - matches!(error, Error::Unsupported(_)), - "{forwarded} declined as {error:?}, which the host would not fall back on" - ); - } -} - -#[test] -fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { - // `get_request_headers` assigns `headers["Authorization"]` unconditionally - // once a bearer token resolves, so the deployment's identity wins on - // Python. Keeping the caller's would authorize and bill the call as a - // different principal, and only when the deployment carries `rust: true`. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.extra_headers = Some(Map::from_iter([( - "Authorization".to_string(), - json!("Bearer caller-supplied"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let authorizations: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.as_str()) - .collect(); - assert_eq!( - authorizations, - vec!["Bearer sk-test"], - "the deployment token must be the only authorization on the wire" - ); -} - -#[test] -fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { - // The opposite precedence, and deliberate: Anthropic's own transform - // honours a forwarded OAuth bearer, so the Bedrock fix above must not be - // generalized into a rule that the configured key always wins. - // - // An OAuth bearer is the whole of that exception. This forwarded a plain - // `x-api-key` until round 17, which read as the same claim and was not: - // Python overwrites a forwarded `x-api-key` with the deployment's. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "authorization".to_string(), - json!("Bearer sk-ant-oat01-forwarded"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .map(|(_, value)| value.as_str()) - .collect(); - assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-forwarded") - ); -} - -#[test] -fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { - // The configured bearer identity has its own account and quota boundary, - // so a request carrying one must not be signed as whatever principal the - // host's AWS credentials resolve to. - let prepared = prepare_chat_completions_call(request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - )) - .expect("prepares"); - assert_eq!( - prepared.auth, - RequestAuth::Bearer { - token: "sk-test".to_string() - } - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-test"), - "prepare did not carry the bearer token" - ); -} - -fn decline_reason( - model: &str, - provider: Option<&str>, - messages: Value, - params: Value, -) -> Option<&'static str> { - let params = match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }; - super::chat_completions_decline_reason(model, provider, messages, ¶ms) -} - -#[test] -fn the_gate_accepts_what_prepare_accepts() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - ), - None - ); -} - -#[test] -fn the_gate_declines_without_resolving_credentials_or_calling_out() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ), - Some("streaming") - ); - assert_eq!( - decline_reason( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!("nope"), - json!({}) - ), - Some("unreadable message list") - ); - assert_eq!( - decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), - Some("empty message list") - ); -} - -#[test] -fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { - // A gate that accepts what prepare then declines would make the host emit - // its pre-call logging on a path that falls back, so pin the agreement. - for (messages, params) in [ - ( - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 8}), - ), - ( - json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), - json!({"temperature": 0.1}), - ), - ( - json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), - json!({}), - ), - ] { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params.clone() - ), - None, - "gate declined {messages}" - ); - prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params, - )) - .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); - } -} - -mod round_trip { - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; - - use super::*; - use crate::chat_completions::chat_completions; - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - fn http_response(status: &str, body: &str) -> String { - format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - } - - /// Serve one request from a stub upstream and hand back what it received. - async fn serve_once( - status: &'static str, - body: &'static str, - ) -> (String, tokio::task::JoinHandle) { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let port = listener.local_addr().expect("addr").port(); - let handle = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts"); - let received = read_http_request(&mut socket).await; - socket - .write_all(http_response(status, body).as_bytes()) - .await - .expect("writes response"); - socket.flush().await.expect("flushes"); - received - }); - (format!("http://127.0.0.1:{port}/v1/messages"), handle) - } - - fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { - ChatCompletionsRequest { - model: "anthropic/claude-sonnet-4-5", - messages, - optional_params: match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: Some(api_base), - custom_llm_provider: None, - extra_headers: None, - timeout: Some(std::time::Duration::from_secs(10)), - } - } - - const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; - - #[tokio::test] - async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { - let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; - let response = chat_completions(call( - &api_base, - json!([ - {"role": "system", "content": "be terse"}, - {"role": "user", "content": "hi"} - ]), - json!({"max_tokens": 16}), - )) - .await - .expect("call succeeds"); - - let received = handle.await.expect("server task"); - let sent: Value = serde_json::from_str( - received - .split_once("\r\n\r\n") - .expect("request has a body") - .1, - ) - .expect("body is json"); - assert_eq!( - sent["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - sent["system"], - json!([{"type": "text", "text": "be terse"}]) - ); - assert_eq!(sent["max_tokens"], json!(16)); - assert!(received.to_lowercase().contains("x-api-key: sk-test")); - - assert_eq!( - response.choices[0].message.content.as_deref(), - Some("hello") - ); - assert_eq!(response.usage.total_tokens, 15); - } - - #[tokio::test] - async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { - // The provider was called and billed, so the host must not retry this - // on its own path. `MissingField` here would read as a pre-send - // decline and be retried; `InvalidResponse` cannot. - const NO_USAGE: &str = - r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; - let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { - const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; - let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn an_upstream_error_status_keeps_its_code() { - let (api_base, handle) = - serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("upstream rejects"); - handle.await.expect("server task"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "expected a 429, got {err:?}" - ); - } - - #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. - let port = { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - listener.local_addr().expect("has an address").port() - // Dropped here, so the port is closed and the connect is refused. - }; - let err = chat_completions(call( - &format!("http://127.0.0.1:{port}/v1/messages"), - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("nothing is listening"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Connect(_)) - ), - "expected a pre-send connect failure, got {err:?}" - ); - } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string() - })), - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); - } -} diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 015e026f6da..4327754ed05 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -26,3 +26,186 @@ pub(super) fn string_headers( ) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) } + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use futures_util::future::BoxFuture; + use litellm_secrets::{SecretValue, source::SecretSource}; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::{messages_provider_config, string_headers, truncate_error_body}; + use crate::messages::{ + Error, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + types::MessagesShaping, + }; + + struct RecordingSecrets { + values: Vec<(&'static str, String)>, + requested: std::sync::Mutex>, + } + + impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + Ok(self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } + } + + fn secrets_call() -> MessagesCall { + let Value::Object(body) = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }) else { + unreachable!("literal object") + }; + MessagesCall { + model: "claude-sonnet-4-5".into(), + body, + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } + } + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + #[tokio::test] + async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + let secrets = Arc::new(RecordingSecrets { + values: vec![ + ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), + ("ANTHROPIC_BASE_URL", format!("http://{addr}")), + ], + requested: std::sync::Mutex::new(Vec::new()), + }); + + let output = litellm_host::run::run( + messages_machine(secrets.clone()), + &LocalMessagesHost::new(secrets_call()), + ) + .await + .expect("messages request succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = server.await.expect("server task completes"); + assert!( + request + .to_ascii_lowercase() + .contains("x-api-key: sk-from-manager"), + "{request}" + ); + let requested = secrets.requested.lock().unwrap().clone(); + assert_eq!( + requested, + messages_provider_config("anthropic") + .unwrap() + .secret_names() + .iter() + .map(ToString::to_string) + .collect::>() + ); + } + + #[test] + fn provider_config_resolves_anthropic_and_azure_ai() { + assert!(messages_provider_config("anthropic").is_some()); + assert!(messages_provider_config("azure_ai").is_some()); + assert!(messages_provider_config("openai").is_none()); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(400); + let truncated = truncate_error_body(&body); + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = json!({"x-count": 3}).as_object().unwrap().clone(); + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + Error::Headers(litellm_http::request::HeaderError { + context: "messages", + name: "x-count".to_string(), + actual: "number", + }) + ); + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 8795d4f8507..180eb08810e 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -46,6 +46,3 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } + } + + /// What the host does to the wire request in `before_send`. + #[derive(Clone, Copy, Debug)] + enum Host { + Detached, + ReplacesDocument, + } + + const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + + impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } + } + + struct Sent { + result: Result<(), Error>, + provider_body: Option, + } + + async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = + mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } + } + + fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) + } + + #[rstest] + #[case::azure_ai(Route::AzureAi)] + #[case::vertex_mistral(Route::VertexMistral)] + #[case::azure_cohere_parse(Route::AzureCohereParse)] + #[tokio::test] + async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); + } + + #[rstest] + #[tokio::test] + async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, + ) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index f298f106a5f..270a402c9fa 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -9,36 +9,210 @@ pub mod types; pub mod wire; #[cfg(test)] -#[path = "../../tests/aws_textract_ocr.rs"] -mod aws_textract_tests; +pub(crate) mod test_support { + use std::sync::{Arc, Mutex}; -#[cfg(test)] -#[path = "../../tests/azure_ai_ocr.rs"] -mod azure_ai_tests; -#[cfg(test)] -#[path = "../../tests/azure_document_intelligence_ocr.rs"] -mod azure_document_intelligence_tests; -#[cfg(test)] -#[path = "../../tests/cohere_ocr.rs"] -mod cohere_tests; -#[cfg(test)] -#[path = "../../tests/deepseek_ocr.rs"] -mod deepseek_tests; -#[cfg(test)] -#[path = "../../tests/ocr/document.rs"] -mod document_tests; -#[cfg(test)] -#[path = "../../tests/reducto_ocr.rs"] -mod reducto_tests; -#[cfg(test)] -#[path = "../../tests/ocr/support.rs"] -pub(crate) mod test_support; -#[cfg(test)] -#[path = "../../tests/ocr.rs"] -pub(crate) mod tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_deepseek_ocr.rs"] -mod vertex_ai_deepseek_tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_ocr.rs"] -mod vertex_ai_tests; + use futures_util::future::BoxFuture; + use litellm_host::event::WireRequest; + use litellm_llms::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::LiteLLMOcrResponse, + }; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + use crate::ocr::{ + route::{LocalOcrHost, ocr_machine}, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, + }; + + /// Stands in for a host with no hooks registered: the wire request goes out unchanged + /// and response events go nowhere. + pub(crate) struct NoHooks; + + impl CallHooks for NoHooks { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(wire) }) + } + + fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + } + + pub(crate) fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) + } + + pub(crate) async fn perform_ocr( + request: LiteLLMOcrRequest, + ) -> Result { + crate::ocr::client::perform(&ocr_client(), request).await + } + + pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { + litellm_host::run::run(ocr_machine(ocr_client()), &host).await + } + + pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + base, + json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + options, + ) + } + + pub(crate) fn wire_request_with_document( + model: &str, + base: &str, + document: Value, + options: Value, + ) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() + } + + pub(crate) fn resolved_request( + request: LiteLLMOcrRequest, + ) -> crate::ocr::types::ResolvedOcrRequest { + request + .map_document(crate::ocr::document::prepare_document) + .unwrap() + } + + pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let request = resolved_request(request); + let document = request.document.clone().with_source(source.into()); + request.with_document(document.into()) + } + + pub(crate) fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; + + /// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. + pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buffer = [0u8; 4096]; + let _ = socket.read(&mut buffer).await.unwrap(); + let head = format!( + "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + SERVED_DOCUMENT.len() + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(SERVED_DOCUMENT).await.unwrap(); + } + }); + (base, task) + } + + pub(crate) struct MockResponse { + pub status: u16, + pub headers: Vec<(&'static str, String)>, + pub body: Value, + } + + impl MockResponse { + pub fn json(body: Value) -> Self { + Self { + status: 200, + headers: vec![], + body, + } + } + } + + pub(crate) async fn mock_server( + responses: Vec, + ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = requests.clone(); + let server_base = base.clone(); + let task = tokio::spawn(async move { + for response in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + let mut buffer = [0u8; 4096]; + let header_end = loop { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { + break index + 4; + } + }; + let length = String::from_utf8_lossy(&bytes[..header_end]) + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + while bytes.len() < header_end + length { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&bytes).into_owned()); + let body = serde_json::to_vec(&response.body).unwrap(); + let headers = response + .headers + .into_iter() + .map(|(name, value)| { + format!("{name}: {}\r\n", value.replace("{base}", &server_base)) + }) + .collect::(); + let head = format!( + "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", + response.status, + body.len(), + headers + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(&body).await.unwrap(); + } + }); + (base, requests, task) + } + + pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { + request + .lines() + .take_while(|line| !line.is_empty()) + .find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case(name).then(|| value.trim()) + }) + } +} diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index 26c9ac27102..7f83291bdab 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -207,3 +207,3589 @@ impl litellm_host::host::Host for LocalOcrHost { Ok(()) } } + +#[cfg(test)] +mod aws_textract_tests { + use std::{collections::BTreeMap, time::SystemTime}; + + use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + use time::{PrimitiveDateTime, format_description}; + + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, header, mock_server, perform_ocr_with, request_body, + wire_request_with_document, + }, + types::LiteLLMOcrRequest, + }; + + const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; + const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; + + fn textract_request(base: &str) -> LiteLLMOcrRequest { + textract_request_for("aws_textract/detect-document-text", base) + } + + fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) + } + + fn textract_response() -> MockResponse { + MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) + } + + /// Recomputes SigV4 over the bytes the server received, at the time the client claimed. + fn expected_authorization(url: &str, raw_request: &str) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| { + ( + name.to_string(), + header(raw_request, name).unwrap().to_string(), + ) + }) + .collect(); + let body = raw_request.split_once("\r\n\r\n").unwrap().1; + sign_post( + url, + body.as_bytes(), + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() + } + + #[tokio::test] + async fn the_request_is_signed_for_textract_and_lines_become_the_page() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + + let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + header(&raw, "content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); + } + + #[tokio::test] + async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + } + + #[tokio::test] + async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { + let (base, _, server) = mock_server(vec![MockResponse { + status: 400, + headers: vec![], + body: json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + }]) + .await; + + let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap_err(); + server.await.unwrap(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); + } + + #[tokio::test] + async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + let request = textract_request_for("aws_textract/analyze-document", &base); + + let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!( + request_body(&raw)["FeatureTypes"], + json!(["LAYOUT", "TABLES"]) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); + } +} + +#[cfg(test)] +mod azure_ai_tests { + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + #[tokio::test] + async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); + } + + #[tokio::test] + async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.credentials.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); + } + + #[tokio::test] + async fn rejects_non_inline_body_after_guardrails() { + let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { + wire.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(wire) + }); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); + } + + mod transformation { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use litellm_auth::{ + ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, + }; + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::ocr::{ + test_support::{MockResponse, header, mock_server, perform_ocr}, + types::LiteLLMOcrRequest, + wire::decode_request, + }; + + #[derive(Debug)] + struct CountingToken { + token: fn(usize) -> String, + calls: AtomicUsize, + } + + impl CountingToken { + fn new(token: fn(usize) -> String) -> Arc { + Arc::new(Self { + token, + calls: AtomicUsize::new(0), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } + } + + impl TokenProvider for CountingToken { + fn acquire(&self) -> TokenFuture<'_> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let token = SecretValue::new((self.token)(call)); + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token, + expires_on: None, + }) + }) + } + } + + fn numbered_token(call: usize) -> String { + format!("callback-{call}") + } + + fn azure_request( + provider: &Arc, + api_base: Option<&str>, + api_key: Option<&str>, + extra_headers: Value, + optional_params: Value, + ) -> LiteLLMOcrRequest { + let wire = serde_json::from_value(json!({ + "model": "azure_ai/mistral-ocr-latest", + "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": null, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": 2.0 + })) + .unwrap(); + LiteLLMOcrRequest { + azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), + ..decode_request(wire).unwrap() + } + } + + fn ocr_page() -> MockResponse { + MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) + } + + #[tokio::test] + async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; + + for _ in 0..2 { + perform_ocr(azure_request( + &provider, + Some(&base), + None, + Value::Null, + json!({}), + )) + .await + .unwrap(); + } + server.await.unwrap(); + + assert_eq!(provider.calls(), 2); + let requests = seen.lock().unwrap(); + assert_eq!( + requests + .iter() + .map(|request| header(request, "authorization")) + .collect::>(), + [Some("Bearer callback-1"), Some("Bearer callback-2")] + ); + } + + #[rstest] + #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] + #[case::provider_beats_static_token( + None, + Value::Null, + json!({"azure_ad_token":"static-token"}), + "Bearer callback-1", + 1 + )] + #[case::header_wins_on_the_wire_but_provider_still_runs( + None, + json!({"Authorization":"Bearer override"}), + json!({}), + "Bearer override", + 1 + )] + #[tokio::test] + async fn credential_precedence( + #[case] api_key: Option<&str>, + #[case] extra_headers: Value, + #[case] optional_params: Value, + #[case] expected_authorization: &str, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + perform_ocr(azure_request( + &provider, + Some(&base), + api_key, + extra_headers, + optional_params, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(provider.calls(), expected_calls); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + header(&requests[0], "authorization"), + Some(expected_authorization) + ); + } + + #[rstest] + #[case::missing_api_base( + false, + json!({}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: "AZURE_AI_API_BASE", + })), + 0 + )] + #[case::unsupported_oidc_reference( + true, + json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), + 0 + )] + #[case::empty_provider_token_ignores_static_token( + true, + json!({"azure_ad_token":"static-token"}), + |_| String::new(), + |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + 1 + )] + #[tokio::test] + async fn credential_failures_send_no_provider_request( + #[case] with_api_base: bool, + #[case] optional_params: Value, + #[case] token: fn(usize) -> String, + #[case] expected: fn(&Error) -> bool, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + let error = perform_ocr(azure_request( + &provider, + with_api_base.then_some(base.as_str()), + None, + Value::Null, + optional_params, + )) + .await + .unwrap_err(); + server.abort(); + + assert!(expected(&error), "unexpected error: {error:?}"); + assert_eq!(provider.calls(), expected_calls); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod azure_document_intelligence_tests { + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + fn query_value(url: &str, key: &str) -> Option { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find_map(|(name, value)| (name == key).then(|| value.into_owned())) + } + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), + ); + request.document = serde_json::from_value::< + litellm_llms::base_llm::ocr::transformation::OcrDocument, + >(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf"}) + ); + } + + #[rstest] + #[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] + #[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] + #[case(json!({"features":"languages&pages=1"}), Error::Features)] + #[case(json!({"req_format":"azure"}), Error::RequestFormat)] + #[tokio::test] + async fn rejects_invalid_pages_features_and_format( + #[case] options: Value, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: Some(base), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }); + let result = match result { + Ok(request) => perform_ocr(request).await, + Err(error) => Err(error), + }; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[rstest] + #[case(json!({}))] + #[case(json!({"req_format":"litellm"}))] + #[tokio::test] + async fn missing_native_fields_keep_page_text_without_retaining_raw_response( + #[case] options: Value, + ) { + let operation = json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} + }); + let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; + let response = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + options, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + assert_eq!(serialized.get("content"), Some(&Value::Null)); + assert_eq!(serialized.get("tables"), Some(&Value::Null)); + assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + let target = requests[0].split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + for field in ["pages", "features", "req_format"] { + assert_eq!(query_value(&url, field), None); + } + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn inline_document_decodes_to_base64_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + } + + #[tokio::test] + async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = crate::ocr::client::perform( + &client, + wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + + let target = seen.lock().unwrap()[0] + .split_whitespace() + .nth(1) + .unwrap() + .to_string(); + assert_eq!( + query_value(&format!("{base}{target}"), "api-version").as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":612,"height":792,"dpi":72}) + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { + return; + }; + match request_count.lock().unwrap().len() { + 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), + 2 => assert!(raw.body.contains("succeeded")), + count => panic!("unexpected callback after {count} requests"), + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn polling_forwards_bearer_credentials() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer token") + ); + } + + #[tokio::test] + async fn polling_does_not_follow_redirects() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 302, + headers: vec![("Location", "{base}/redirected".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(seen.lock().unwrap().len(), 2); + server.abort(); + } + + #[tokio::test] + async fn polling_rejects_terminal_failure() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"failed"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("status failed")); + } + + #[tokio::test] + async fn malformed_provider_pages_report_response_paths() { + for (analysis, path) in [ + (json!({"pages":null}), "pages"), + (json!({"pages":[null]}), "pages[0]"), + (json!({"pages":[{"lines":null}]}), "lines"), + (json!({"pages":[{"width":"bad"}]}), "width"), + ] { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":analysis + }))]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains(path), "{error}"); + } + } + + #[tokio::test] + async fn rejects_missing_invalid_and_cross_origin_operation_locations() { + for headers in [ + Vec::new(), + vec![("Operation-Location", "/relative".into())], + vec![("Operation-Location", "http://example.com/operation".into())], + vec![( + "Operation-Location", + "http://user:password@127.0.0.1/operation".into(), + )], + ] { + let (base, _, server) = mock_server(vec![MockResponse { + status: 202, + headers, + body: json!({}), + }]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("operation-location")); + } + } + + #[tokio::test] + async fn polling_deadline_bounds_retry_delay() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "9999".into())], + body: json!({"status":"notStarted"}), + }, + ]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: std::time::Duration::from_millis(100), + ..OcrSettings::default() + }); + + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + crate::ocr::client::perform(&client, request), + ) + .await + .unwrap() + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("timed out")); + } + + #[tokio::test] + async fn model_id_is_encoded_and_dot_segments_are_rejected() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + perform_ocr(wire_request( + "azure_ai/doc-intelligence/a ?#é", + &base, + json!({}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); + + for model in [ + "azure_ai/doc-intelligence/.", + "azure_ai/doc-intelligence/..", + ] { + let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) + .await + .unwrap_err(); + assert!(error.to_string().contains("dot segment")); + } + } + + mod transformation { + use std::sync::{Arc, Mutex}; + + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::transformation::OcrDocument; + use serde_json::{Value, json}; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), + ); + request.document = serde_json::from_value::(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) + ); + } + + #[tokio::test] + async fn rejects_invalid_pages_features_and_format() { + for options in [ + json!({"pages":[true]}), + json!({"pages":[1,"2"]}), + json!({"pages":[-1]}), + json!({"pages":"1&&features=bad"}), + json!({"features":"languages&pages=1"}), + json!({"req_format":"azure"}), + ] { + let request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + "http://127.0.0.1:1", + options.clone(), + ); + let rejected = perform_ocr(request).await.is_err(); + assert!(rejected, "accepted {options}"); + } + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let request_count = seen.clone(); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed + .lock() + .unwrap() + .push((request_count.lock().unwrap().len(), raw.body.clone())); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + assert_eq!( + *responses_received.lock().unwrap(), + [ + (1, r#"{"submitted":true}"#.to_string()), + (2, r#"{"status":"succeeded"}"#.to_string()), + ] + ); + } + } +} + +#[cfg(test)] +mod cohere_tests { + mod transformation { + use litellm_llms::{ + base_llm::ocr::{ + error::Error, + transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, + }, + cohere::ocr::transformation::*, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + #[tokio::test] + async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({ + "output_format":"markdown", "timeout":30, + "extra_body":{ + "output_format": {"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + } + }), + ); + let request = request.with_document( + serde_json::from_value(json!({ + "type":"image_url","image_url":"https://example.com/original.png" + })) + .unwrap(), + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model":"parse", "output_format":{"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({"output_format":null,"req_format":null}), + ); + let request = request.with_document( + serde_json::from_value( + json!({"type":"image_url","image_url":"https://example.com/a.png"}), + ) + .unwrap(), + ); + assert_eq!( + request.response_format().unwrap(), + OcrResponseFormat::Litellm + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[rstest] + #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] + #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] + #[tokio::test] + async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( + #[case] model: &str, + #[case] request_line: &str, + ) { + use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = crate::ocr::test_support::wire_request(model, &base, json!({})) + .with_document( + serde_json::from_value::( + json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), + ) + .unwrap() + .into(), + ); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with(request_line), "{}", requests[0]); + assert_eq!( + header(&requests[0], "authorization"), + Some("Bearer test-key") + ); + } + + #[rstest] + #[tokio::test] + async fn route_rejects_non_image_document_without_a_request( + #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, + ) { + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + + let error = perform_ocr(crate::ocr::test_support::wire_request( + model, + &base, + json!({}), + )) + .await + .unwrap_err(); + server.abort(); + + assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod deepseek_tests { + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, + vertex_ai::ocr::deepseek_transformation::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, + normalize_response as transform_ocr_response, + }, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + fn document() -> OcrDocument { + serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() + } + + #[rstest] + #[case("stream", json!(true))] + #[case("temperature", json!(0.1))] + #[case("max_tokens", json!(1024))] + #[case("top_p", json!(0.9))] + #[case("n", json!(2))] + #[case("stop", json!("done"))] + #[case("stop", json!(["done", "stop"]))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: DeepSeekOcrParams = + serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); + let result = serde_json::to_value( + VertexAIDeepSeekOCRConfig + .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); + } + + #[rstest] + #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] + #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] + fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + &[], + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); + } + + #[rstest] + #[case(json!("# hello"), "# hello")] + #[case(json!("{broken"), "{broken")] + #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] + #[case(json!({"pages":[]}), "")] + #[case(json!("[]"), "[]")] + #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] + #[case(json!({"pages":[{"markdown":"object"}]}), "object")] + fn response_codec_handles_text_json_and_objects( + #[case] content: Value, + #[case] expected: &str, + ) { + let structured = content + .as_object() + .is_some_and(|object| object.contains_key("pages")) + || content + .as_str() + .is_some_and(|text| text.contains("\"pages\"")); + let response: DeepSeekOcrResponse = serde_json::from_value( + json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), + ) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["markdown"], expected); + assert_eq!(result["pages"][0]["index"], 0); + if structured { + assert!(result["usage_info"].is_null()); + } else { + assert_eq!(result["usage_info"]["prompt_tokens"], 1); + } + } + + #[test] + fn structured_result_maps_pages_usage_model_and_annotation() { + let response: DeepSeekOcrResponse = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], + "model":"provider-model", + "usage_info":{"pages_processed":1}, + "document_annotation":{"language":"en"}, + "future":"kept" + }}}] + })) + .unwrap(); + let result = transform_ocr_response("requested", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["index"], 2); + assert_eq!(result["pages"][0]["images"][0]["id"], "one"); + assert_eq!(result["model"], "provider-model"); + assert_eq!(result["usage_info"]["pages_processed"], 1); + assert_eq!(result["document_annotation"]["language"], "en"); + assert_eq!(result["future"], "kept"); + } + + #[test] + fn response_codec_rejects_missing_empty_and_malformed_content() { + for value in [ + json!({"choices":[{"message":{"content":{}}}]}), + json!({"choices":[]}), + json!({"choices":[{"message":{"content":""}}]}), + json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), + json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), + ] { + let result = serde_json::from_value::(value) + .map_err(|_| ()) + .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } + } +} + +#[cfg(test)] +mod reducto_tests { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers( + #[case] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, + ) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let document = if mime_type.starts_with("image/") { + json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) + } else { + json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) + }; + let mut request = crate::ocr::types::LiteLLMOcrRequest { + document: serde_json::from_value::(document) + .unwrap() + .into(), + ..wire_request(&format!("reducto/{model}"), &base, json!({})) + }; + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; + assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); + assert!(multipart.contains("\r\n\r\nabc\r\n--")); + assert!(requests[1].starts_with("POST /parse ")); + let source_field = if model == "parse-legacy" { + "document_url" + } else { + "input" + }; + assert_eq!( + request_body(&requests[1]), + json!({source_field:"reducto://uploaded.pdf"}) + ); + for request in requests.iter() { + assert!( + request + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + } + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case(json!({"file_id":""}))] + #[case(json!({}))] + #[case(json!({"file_id":null}))] + #[tokio::test] + async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("file_id")); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn upload_failure_stops_before_parse() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 503, + headers: vec![], + body: json!({"error":"unavailable"}), + }]) + .await; + assert!( + perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .is_err() + ); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[rstest] + #[case("https://example.com/a.pdf", Error::ReductoSource)] + #[case("reducto://", Error::RequestField { path: "document file id".into() })] + #[case("data:application/pdf;base64", Error::InvalidDataUri)] + #[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network( + #[case] source: &str, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + source, + ); + let result = perform_ocr(request).await; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[test] + fn response_normalization_groups_blocks_and_distinguishes_null_result() { + use litellm_llms::reducto::ocr::transformation::{ + ReductoResponse, normalize_response as transform_ocr_response, + }; + + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} + ]}}); + let response: ReductoResponse = serde_json::from_value(raw).unwrap(); + let normalized = transform_ocr_response("parse-v3", response) + .unwrap() + .into_json(); + assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); + assert_eq!(normalized["pages"][1]["markdown"], "B"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); + assert_eq!(normalized["usage_info"]["pages_processed"], 2); + assert_eq!(normalized["usage_info"]["credits"], 3.0); + + let missing: ReductoResponse = + serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); + let missing = transform_ocr_response("parse-v3", missing).unwrap(); + assert_eq!(missing.pages[0].markdown, "text"); + let null: ReductoResponse = serde_json::from_value( + json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), + ) + .unwrap(); + let null = transform_ocr_response("parse-v3", null).unwrap(); + assert!(null.pages.is_empty()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[tokio::test] + async fn native_format_retains_the_provider_response() { + let raw = json!({ + "result":{"chunks":[{"content":"native OCR response"}]}, + "usage":{"num_pages":1} + }); + let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages[0].markdown, "native OCR response"); + assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); + } + + #[tokio::test] + async fn unknown_model_reaches_parse_and_keeps_its_name() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[{"content":"future model response"}]} + }))]) + .await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/future-parse-model", &base, json!({})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.model, "future-parse-model"); + assert_eq!(response.pages[0].markdown, "future model response"); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!( + request_body(&requests[0]), + json!({"input":"reducto://ready.pdf"}) + ); + } + + #[tokio::test] + async fn guardrail_rewrites_document_before_upload() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_before_send(|wire, _| { + assert_eq!( + wire.body["document_url"], + "data:application/pdf;base64,YWJj" + ); + Ok(WireRequest { + body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert!(requests[0].contains("reducto://guarded.pdf")); + } + + mod transformation { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, + reducto::ocr::transformation::*, + }; + use rstest::rstest; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn v3_options_preserve_explicit_null() { + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = crate::ocr::test_support::ocr_client(); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({ + "input":"reducto://ready.pdf", "formatting":null, "settings":{} + }) + ); + let absent = ReductoParseV3Config + .map_ocr_params( + &litellm_core_utils::call_arguments::CallArguments::default(), + "parse-v3", + ) + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + assert!(requests[0].contains("application/pdf")); + assert!(requests[0].contains("abc")); + assert!(requests[1].starts_with("POST /parse ")); + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case("https://example.com/a.pdf")] + #[case("reducto://")] + #[case("data:application/pdf;base64")] + #[case("data:application/pdf;base64,INVALID!")] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + source, + ); + assert!(perform_ocr(request).await.is_err()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[rstest] + #[case("reducto/parse-v3")] + #[case("reducto/parse-legacy")] + #[tokio::test] + async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let mut request = wire_request(model, &base, json!({})); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer original".into())]; + let host = LocalOcrHost::new(request).with_before_send(|wire, _| { + Ok(WireRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!(requests[1].starts_with("POST /parse ")); + for request in requests.iter() { + assert!(request.contains("authorization: Bearer guarded")); + assert!(!request.contains("Bearer original")); + } + } + } +} + +#[cfg(test)] +mod vertex_ai_tests { + use litellm_auth::InputSource; + use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); + } + + #[tokio::test] + async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + crate::ocr::client::perform( + &client, + wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].starts_with( + "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + } + + #[tokio::test] + async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); + } + + #[tokio::test] + async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/mistral-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + #[tokio::test] + async fn adapters_build_complete_requests_and_share_mistral_normalization() { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "ignored" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "ignored" + }) + ); + } + let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); + let raw = serde_json::to_vec(&payload).unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + + mod transformation { + + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::test_support::wire_request; + + #[rstest] + #[case::mistral(false)] + #[case::vertex(true)] + #[tokio::test] + async fn configs_build_complete_requests_and_share_mistral_normalization( + #[case] use_vertex: bool, + ) { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + let http = if use_vertex { + &vertex_http + } else { + &direct_http + }; + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "preserved" + }) + ); + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &payload, Default::default()) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &payload, Default::default()) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + } +} + +#[cfg(test)] +mod vertex_ai_deepseek_tests { + use litellm_auth::InputSource; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + + #[test] + fn host_registration_selects_deepseek_without_affecting_mistral() { + assert!(crate::ocr::arguments::is_supported_request( + "deepseek-ocr-maas", + Some("vertex_ai") + )); + assert!(crate::ocr::arguments::is_supported_request( + "mistral-ocr-maas", + Some("vertex_ai") + )); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/deepseek-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + mod deepseek_transformation { + use serde_json::json; + + use super::*; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = + crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + } +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::{Arc, Mutex}; + + use futures_util::future::BoxFuture; + use litellm_auth_gcp::VertexAuth; + use litellm_host::{ + event::{CallEvent, MachineEvent, WireRequest}, + host::{Host, HostOp, HostResult}, + machine::{HostFailure, Machine, MachineStep}, + }; + use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, + }; + use litellm_llms::base_llm::ocr::{ + error::Error as OcrError, + handler::OcrClient, + settings::OcrSettings, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, + }, + }; + use litellm_secrets::source::SecretSource; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + struct RecordingSecretSource { + names: Arc>>, + values: &'static [(&'static str, &'static str)], + api_base: String, + } + + impl SecretSource for RecordingSecretSource { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> + { + self.names.lock().unwrap().push(name.to_owned()); + Box::pin(async move { + Ok(match name { + "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), + _ => self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + } + .map(litellm_secrets::SecretValue::new)) + }) + } + } + + #[rstest] + #[case::mistral("mistral/model", json!({}))] + #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] + #[tokio::test] + async fn ocr_contract_upstream_error_preserves_status_body_and_headers( + #[case] model: &str, + #[case] options: Value, + ) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let (base, seen, server) = mock_server(vec![MockResponse { + status: 422, + headers: vec![ + ("Retry-After", "17".into()), + ("X-Request-ID", "request-123".into()), + ("X-Future-Header", "retained".into()), + ], + body: payload, + }]) + .await; + let error = perform_ocr(wire_request(model, &base, options)) + .await + .unwrap_err(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + let OcrError::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) + ); + } + assert_eq!( + body.len(), + expected_body.len(), + "provider error body was truncated" + ); + assert_eq!(body, expected_body); + } + + #[test] + fn request_boundary_selects_mistral_and_rejects_unknown_providers() { + let request = OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"extract_header":true,"unknown":42}) + .as_object() + .unwrap() + .clone(), + input_sources: Default::default(), + timeout_seconds: None, + }; + assert!(decode_request(request).is_ok()); + assert!( + decode_request(OcrWireRequest { + model: "model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: Some("unknown".into()), + extra_headers: None, + optional_params: serde_json::Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) + .is_err() + ); + } + + #[tokio::test] + async fn facade_executes_direct_mistral_once() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let result = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", + "extract_header":true, + "unknown":"ignored" + }) + ); + } + + #[tokio::test] + async fn facade_retains_native_response_when_requested() { + let provider_response = json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1}, + "provider_only":"preserved" + }); + let (base, _, server) = + mock_server(vec![MockResponse::json(provider_response.clone())]).await; + let response = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(provider_response) + ); + } + + #[rstest] + #[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] + #[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] + #[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] + #[tokio::test] + async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected_key: &str, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: secrets, + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); + } + + #[tokio::test] + async fn mistral_ocr_resolves_provider_secrets_before_transformation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: &[("MISTRAL_API_KEY", "source-key")], + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); + } + + #[tokio::test] + async fn ocr_client_uses_the_injected_http_pool_configuration() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&settings).config, + UrlPolicy::default(), + VertexAuth::default(), + OcrSettings::default(), + Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + ) + .unwrap(); + crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); + } + + fn event_name(event: &CallEvent) -> &'static str { + match event { + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", + CallEvent::Succeeded { .. } => "success", + CallEvent::Failed { .. } => "failure", + } + } + + fn recording_host( + request: crate::ocr::types::LiteLLMOcrRequest, + events: Arc>>, + block: bool, + ) -> LocalOcrHost { + let before_send_events = events.clone(); + LocalOcrHost::new(request) + .with_before_send(move |wire, _| { + before_send_events.lock().unwrap().push("before_send"); + if block { + return Err(OcrError::InvalidRequest("blocked".into())); + } + Ok(wire) + }) + .with_observer(move |event| events.lock().unwrap().push(event_name(event))) + } + + #[tokio::test] + async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_before_send(|mut wire, _| { + wire.headers + .push(("x-core-callback".into(), "edited".into())); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); + } + + #[tokio::test] + async fn before_send_context_names_the_route_and_its_secrets() { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let host = LocalOcrHost::new(wire_request( + "mistral/model", + &base, + json!({"pages": [0], "req_format": "native"}), + )) + .with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some((wire.clone(), context.clone())); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let (wire, context) = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.custom_llm_provider, "mistral"); + assert_eq!(context.model, "model"); + assert_eq!(wire.body["pages"], json!([0])); + assert!(context.secret_fields.is_empty()); + assert_eq!(context.optional_params["req_format"], "native"); + + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let request = wire_request( + "azure_ai/model", + &base, + json!({"client_secret": "shh", "tenant_id": "t"}), + ); + let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + }); + let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some(context.clone()); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let context = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.secret_fields, ["client_secret"]); + } + + #[tokio::test] + async fn lifecycle_orders_hooks_and_emits_one_success() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "response", "success"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", "http://127.0.0.1:1", json!({})), + events.clone(), + true, + ); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + } + + #[tokio::test] + async fn upstream_failure_emits_one_terminal_failure() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 500, + headers: vec![], + body: json!({"error":"failed"}), + }]) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + assert!(perform_ocr_with(host).await.is_err()); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + /// Drives the machine by hand, answering every op through `host` except `before_send`, + /// which `intercept` answers so a test can fail or cancel exactly there. + async fn drive_until( + client: OcrClient, + host: &LocalOcrHost, + mut intercept: impl FnMut(WireRequest) -> Result>, + ) -> ( + Result, + Vec<&'static str>, + crate::ocr::route::OcrMachine, + ) { + let mut machine = ocr_machine(client); + let mut result = None; + let mut ops = Vec::new(); + let outcome = loop { + let op = match machine.resume(result.take()).await { + Ok(MachineStep::Host(op)) => op, + Ok(MachineStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + }; + let answer = match op { + HostOp::Route(op) => { + ops.push(match op { + OcrOp::ProjectRequest => "ProjectRequest", + OcrOp::ReadDocument => "ReadDocument", + OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", + }); + host.route(op) + .await + .map(HostResult::Route) + .map_err(HostFailure::Error) + } + HostOp::BeforeSend { wire, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) + } + HostOp::Emit(event) => { + let event = CallEvent::Machine(event); + ops.push(event_name(&event)); + host.emit(&event) + .await + .map(|()| HostResult::Emitted) + .map_err(HostFailure::Error) + } + }; + match answer { + Ok(answer) => result = Some(answer), + Err(failure) => break machine.interrupt(failure).await, + } + }; + (outcome, ops, machine) + } + + #[tokio::test] + async fn failed_before_send_does_not_replay_or_reach_transport() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Error(OcrError::InvalidRequest( + "before_send failed".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(None).await.is_err()); + } + + #[tokio::test] + async fn invalid_provider_response_emits_response_received_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed.lock().unwrap().push(raw.body.clone()); + } + }); + let error = perform_ocr_with(host).await.unwrap_err(); + server.await.unwrap(); + assert!(matches!(error, OcrError::ResponseField { .. })); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + *responses_received.lock().unwrap(), + [r#"{"pages":"invalid"}"#] + ); + } + + #[tokio::test] + async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; + server.await.unwrap(); + assert_eq!(outcome.unwrap().pages[0].markdown, "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); + assert!(matches!( + machine.resume(None).await, + Err(OcrError::InvalidRequest(_)) + )); + } + + async fn drive_native_file_call( + request: crate::ocr::types::LiteLLMOcrRequest, + content: Result, + ) -> (Result, usize) { + let reads = Arc::new(Mutex::new(0)); + let counted = reads.clone(); + let content = Mutex::new(Some(content)); + let host = LocalOcrHost::new(request).with_reader(move || { + *counted.lock().unwrap() += 1; + content.lock().unwrap().take().unwrap() + }); + let outcome = perform_ocr_with(host).await; + let reads = *reads.lock().unwrap(); + (outcome, reads) + } + + #[tokio::test] + async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"file"}] + }))]) + .await; + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::HostReader { + mime_type: Some("application/pdf".into()), + }, + ); + let (response, reads) = drive_native_file_call( + request, + Ok(crate::ocr::types::OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + }), + ) + .await; + server.await.unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "file"); + assert_eq!(reads, 1); + assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); + } + + #[tokio::test] + async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let failure = OcrError::InvalidRequest("reader exploded".into()); + let (response, reads) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Err(failure.clone()), + ) + .await; + assert!( + matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") + ); + assert_eq!(reads, 1); + + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Ok(crate::ocr::types::OcrFileContent { + bytes: Default::default(), + file_name: None, + }), + ) + .await; + assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn path_documents_are_read_by_core_without_a_host_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + let (response, reads) = + drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; + server.await.unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "path"); + assert_eq!(reads, 0); + assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); + + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(OcrError::InvalidRequest("unused".into())), + ) + .await; + assert!(matches!( + response.unwrap_err(), + OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound + )); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Cancelled(OcrError::InvalidRequest( + "cancelled".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); + } + + #[tokio::test] + async fn missing_host_result_preserves_pending_operation() { + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); + let mut machine = ocr_machine(ocr_client()); + assert!(matches!( + machine.resume(None).await.unwrap(), + MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) + )); + assert!(machine.resume(None).await.is_err()); + assert!(matches!( + machine + .resume(Some(HostResult::Route(OcrOpResult::Request { + request: Box::new(request), + caller_token: false, + }))) + .await + .unwrap(), + MachineStep::Host(HostOp::BeforeSend { .. }) + )); + } + + async fn read_bounded_response( + response: Vec, + limit: usize, + ) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") + } + + #[tokio::test] + async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + use litellm_llms::base_llm::ocr::error::Error; + + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(Error::TooLarge { limit: 8 }) + )); + } + } + + #[rstest] + #[case::declared("Content-Length: 1000000")] + #[case::chunked("Transfer-Encoding: chunked")] + #[tokio::test] + async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( + #[case] headers: &str, + ) { + let prefix = "x".repeat(4096); + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), prefix.len()) + .await + .unwrap_err(); + match error { + OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { + assert_eq!(status, 429); + assert_eq!(body, prefix); + } + error => panic!("unexpected error: {error}"), + } + } + + #[test] + fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.transport.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } + } + + #[derive(Debug)] + struct PendingToken { + entered: Arc, + dropped: Arc, + } + + struct TokenFutureDrop(Arc); + + impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } + } + + impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } + } + + #[tokio::test] + async fn interrupt_drops_provider_captures_before_returning() { + use std::sync::atomic::{AtomicBool, Ordering}; + + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = crate::ocr::types::LiteLLMOcrRequest { + transport: OcrTransportConfig { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.transport + }, + azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let host = LocalOcrHost::new(request); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { + HostResult::BeforeSend(wire) + } + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = OcrError::InvalidRequest("cancelled".into()); + let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); + assert!( + dropped.load(Ordering::SeqCst), + "interrupt returned while provider captures were still alive" + ); + assert!( + matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + } + + struct CallerTokenHost { + request: Mutex>, + trace: Mutex>, + } + + impl Host for CallerTokenHost { + async fn route(&self, op: OcrOp) -> Result { + match op { + OcrOp::ProjectRequest => { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrOpResult::Request { + request: Box::new(self.request.lock().unwrap().take().unwrap()), + caller_token: true, + }) + } + OcrOp::AcquireAzureAdToken => { + self.trace.lock().unwrap().push("token".into()); + Ok(OcrOpResult::AzureAdToken( + litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( + "caller-token", + )), + )) + } + OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), + } + } + + async fn before_send( + &self, + wire: WireRequest, + _: &litellm_host::event::RequestContext, + ) -> Result { + let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); + let authorization = wire + .headers + .iter() + .find(|(name, _)| is_authorization(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + self.trace + .lock() + .unwrap() + .push(format!("before_send:{authorization}")); + let headers = wire + .headers + .into_iter() + .map(|(name, value)| match is_authorization(&name) { + true => (name, "Bearer edited".to_string()), + false => (name, value), + }) + .collect(); + Ok(WireRequest { headers, ..wire }) + } + } + + #[tokio::test] + async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request("azure_ai/model", &base, json!({})); + request.credentials.api_key = None; + let host = CallerTokenHost { + request: Mutex::new(Some(request)), + trace: Mutex::new(Vec::new()), + }; + + litellm_host::run::run(ocr_machine(ocr_client()), &host) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!( + *host.trace.lock().unwrap(), + ["project", "token", "before_send:Bearer caller-token"] + ); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer edited\r\n") + ); + } + + #[tokio::test] + async fn interrupting_an_in_flight_provider_request_closes_its_connection() { + use tokio::io::AsyncReadExt; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let received = Arc::new(tokio::sync::Notify::new()); + let server_received = received.clone(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0u8; 4096]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut buffer).await.unwrap(); + request.extend_from_slice(&buffer[..read]); + } + server_received.notify_one(); + loop { + if socket.read(&mut buffer).await.unwrap() == 0 { + break; + } + } + }); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = received.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("the stalled provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + + let cancelled = OcrError::InvalidRequest("cancelled".into()); + assert!( + machine + .interrupt(HostFailure::Cancelled(cancelled)) + .await + .is_err() + ); + tokio::time::timeout(std::time::Duration::from_secs(1), server) + .await + .expect("the provider connection stayed open after the interrupt") + .unwrap(); + } +} diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/tests/audio_transcription.rs similarity index 95% rename from litellm-rust/crates/core/src/audio_transcription/tests.rs rename to litellm-rust/crates/core/tests/audio_transcription.rs index 8ccf7a07a0f..aa5aef0149b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -4,11 +4,9 @@ use std::{ thread, }; +use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; use serde_json::{Map, json}; -use super::audio_transcription; -use crate::audio_transcription::types::AudioTranscriptionRequest; - #[tokio::test] async fn bedrock_request_is_signed_and_contains_audio() { let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); diff --git a/litellm-rust/crates/core/tests/aws_textract_ocr.rs b/litellm-rust/crates/core/tests/aws_textract_ocr.rs deleted file mode 100644 index c536317ad5c..00000000000 --- a/litellm-rust/crates/core/tests/aws_textract_ocr.rs +++ /dev/null @@ -1,193 +0,0 @@ -use std::{collections::BTreeMap, time::SystemTime}; - -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; -use time::{PrimitiveDateTime, format_description}; - -use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, header, mock_server, perform_ocr_with, request_body, - wire_request_with_document, - }, - types::LiteLLMOcrRequest, -}; - -const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; -const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; - -fn textract_request(base: &str) -> LiteLLMOcrRequest { - textract_request_for("aws_textract/detect-document-text", base) -} - -fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - &format!("{base}/"), - json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), - json!({ - "aws_access_key_id": ACCESS_KEY_ID, - "aws_secret_access_key": SECRET_ACCESS_KEY, - "aws_region_name": "eu-west-1" - }), - ) -} - -fn textract_response() -> MockResponse { - MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] - })) -} - -/// Recomputes SigV4 over the bytes the server received, at the time the client claimed. -fn expected_authorization(url: &str, raw_request: &str) -> String { - let format = - format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") - .unwrap(); - let signed_at: SystemTime = - PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) - .unwrap() - .assume_utc() - .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] - .into_iter() - .map(|name| { - ( - name.to_string(), - header(raw_request, name).unwrap().to_string(), - ) - }) - .collect(); - let body = raw_request.split_once("\r\n\r\n").unwrap().1; - sign_post( - url, - body.as_bytes(), - &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, - ) - .unwrap()["Authorization"] - .clone() -} - -#[tokio::test] -async fn the_request_is_signed_for_textract_and_lines_become_the_page() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - - let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.DetectDocumentText") - ); - assert_eq!( - header(&raw, "content-type"), - Some("application/x-amz-json-1.1") - ); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "Invoice 12345"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); -} - -#[tokio::test] -async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { - assert!( - !wire - .headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), - "the hook ran after signing" - ); - wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); -} - -#[tokio::test] -async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { - let (base, _, server) = mock_server(vec![MockResponse { - status: 400, - headers: vec![], - body: json!({ - "__type": "UnsupportedDocumentException", - "Message": "Request has unsupported document format" - }), - }]) - .await; - - let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap_err(); - server.await.unwrap(); - - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); - assert!( - body.contains("multi-page documents are not supported"), - "{body}" - ); -} - -#[tokio::test] -async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [ - {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, - {"Id": "t", "BlockType": "LAYOUT_TITLE", - "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} - ] - }))]) - .await; - let request = textract_request_for("aws_textract/analyze-document", &base); - - let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.AnalyzeDocument") - ); - assert_eq!( - request_body(&raw)["FeatureTypes"], - json!(["LAYOUT", "TABLES"]) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "# Quarterly Report"); -} diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs deleted file mode 100644 index 1492aaaeb11..00000000000 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -#[tokio::test] -async fn facade_executes_azure_mistral_with_prepared_auth() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"include_image_base64":true}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![( - "Authorization".into(), - "Bearer python-prepared-token".into(), - )]; - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "include_image_base64":true - }) - ); -} - -#[tokio::test] -async fn facade_acquires_supplied_entra_token_for_final_request() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"azure_ad_token":"rust-owned-token"}), - ); - request.credentials.api_key = None; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer rust-owned-token\r\n") - ); -} - -#[tokio::test] -async fn rejects_non_inline_body_after_guardrails() { - let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { - wire.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(wire) - }); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); -} - -mod transformation { - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, - }; - use rstest::rstest; - use serde_json::json; - - use super::*; - use crate::ocr::{ - test_support::{MockResponse, header, mock_server, perform_ocr}, - types::LiteLLMOcrRequest, - wire::decode_request, - }; - - #[derive(Debug)] - struct CountingToken { - token: fn(usize) -> String, - calls: AtomicUsize, - } - - impl CountingToken { - fn new(token: fn(usize) -> String) -> Arc { - Arc::new(Self { - token, - calls: AtomicUsize::new(0), - }) - } - - fn calls(&self) -> usize { - self.calls.load(Ordering::SeqCst) - } - } - - impl TokenProvider for CountingToken { - fn acquire(&self) -> TokenFuture<'_> { - let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - let token = SecretValue::new((self.token)(call)); - Box::pin(async move { - Ok(ResolvedCredential::AccessToken { - token, - expires_on: None, - }) - }) - } - } - - fn numbered_token(call: usize) -> String { - format!("callback-{call}") - } - - fn azure_request( - provider: &Arc, - api_base: Option<&str>, - api_key: Option<&str>, - extra_headers: Value, - optional_params: Value, - ) -> LiteLLMOcrRequest { - let wire = serde_json::from_value(json!({ - "model": "azure_ai/mistral-ocr-latest", - "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": null, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": 2.0 - })) - .unwrap(); - LiteLLMOcrRequest { - azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), - ..decode_request(wire).unwrap() - } - } - - fn ocr_page() -> MockResponse { - MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) - } - - #[tokio::test] - async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; - - for _ in 0..2 { - perform_ocr(azure_request( - &provider, - Some(&base), - None, - Value::Null, - json!({}), - )) - .await - .unwrap(); - } - server.await.unwrap(); - - assert_eq!(provider.calls(), 2); - let requests = seen.lock().unwrap(); - assert_eq!( - requests - .iter() - .map(|request| header(request, "authorization")) - .collect::>(), - [Some("Bearer callback-1"), Some("Bearer callback-2")] - ); - } - - #[rstest] - #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] - #[case::provider_beats_static_token( - None, - Value::Null, - json!({"azure_ad_token":"static-token"}), - "Bearer callback-1", - 1 - )] - #[case::header_wins_on_the_wire_but_provider_still_runs( - None, - json!({"Authorization":"Bearer override"}), - json!({}), - "Bearer override", - 1 - )] - #[tokio::test] - async fn credential_precedence( - #[case] api_key: Option<&str>, - #[case] extra_headers: Value, - #[case] optional_params: Value, - #[case] expected_authorization: &str, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - perform_ocr(azure_request( - &provider, - Some(&base), - api_key, - extra_headers, - optional_params, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(provider.calls(), expected_calls); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!( - header(&requests[0], "authorization"), - Some(expected_authorization) - ); - } - - #[rstest] - #[case::missing_api_base( - false, - json!({}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - environment_variable: "AZURE_AI_API_BASE", - })), - 0 - )] - #[case::unsupported_oidc_reference( - true, - json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), - 0 - )] - #[case::empty_provider_token_ignores_static_token( - true, - json!({"azure_ad_token":"static-token"}), - |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), - 1 - )] - #[tokio::test] - async fn credential_failures_send_no_provider_request( - #[case] with_api_base: bool, - #[case] optional_params: Value, - #[case] token: fn(usize) -> String, - #[case] expected: fn(&Error) -> bool, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - let error = perform_ocr(azure_request( - &provider, - with_api_base.then_some(base.as_str()), - None, - Value::Null, - optional_params, - )) - .await - .unwrap_err(); - server.abort(); - - assert!(expected(&error), "unexpected error: {error:?}"); - assert_eq!(provider.calls(), expected_calls); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs deleted file mode 100644 index 6dc9bfa5e7e..00000000000 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ /dev/null @@ -1,712 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent}; -use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::LocalOcrHost; - -fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) -} - -#[tokio::test] -async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), - ); - request.document = - serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf"}) - ); -} - -#[rstest] -#[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] -#[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] -#[case(json!({"features":"languages&pages=1"}), Error::Features)] -#[case(json!({"req_format":"azure"}), Error::RequestFormat)] -#[tokio::test] -async fn rejects_invalid_pages_features_and_format( - #[case] options: Value, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: Some(base), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }); - let result = match result { - Ok(request) => perform_ocr(request).await, - Err(error) => Err(error), - }; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid options: {options}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[rstest] -#[case(json!({}))] -#[case(json!({"req_format":"litellm"}))] -#[tokio::test] -async fn missing_native_fields_keep_page_text_without_retaining_raw_response( - #[case] options: Value, -) { - let operation = json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} - }); - let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; - let response = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - options, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "hello"); - assert_eq!(response.provider_native_response, None); - let serialized = response.into_json(); - assert_eq!(serialized.get("content"), Some(&Value::Null)); - assert_eq!(serialized.get("tables"), Some(&Value::Null)); - assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - let target = requests[0].split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - for field in ["pages", "features", "req_format"] { - assert_eq!(query_value(&url, field), None); - } - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); -} - -#[tokio::test] -async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} - }))]) - .await; - let client = ocr_client().with_settings(OcrSettings { - document_intelligence_api_version: "2099-01-01".into(), - document_intelligence_dpi: 72, - ..OcrSettings::default() - }); - - let result = crate::ocr::client::perform( - &client, - wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - - let target = seen.lock().unwrap()[0] - .split_whitespace() - .nth(1) - .unwrap() - .to_string(); - assert_eq!( - query_value(&format!("{base}{target}"), "api-version").as_deref(), - Some("2099-01-01") - ); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":612,"height":792,"dpi":72}) - ); -} - -#[tokio::test] -async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } -} - -#[tokio::test] -async fn accepted_response_emits_response_received_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { - return; - }; - match request_count.lock().unwrap().len() { - 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), - 2 => assert!(raw.body.contains("succeeded")), - count => panic!("unexpected callback after {count} requests"), - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[tokio::test] -async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); -} - -#[tokio::test] -async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); -} - -#[tokio::test] -async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); -} - -#[tokio::test] -async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } -} - -#[tokio::test] -async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } -} - -#[tokio::test] -async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - let client = ocr_client().with_settings(OcrSettings { - poll_timeout: std::time::Duration::from_millis(100), - ..OcrSettings::default() - }); - - let error = tokio::time::timeout( - std::time::Duration::from_secs(1), - crate::ocr::client::perform(&client, request), - ) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); -} - -#[tokio::test] -async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } -} - -mod transformation { - use std::sync::{Arc, Mutex}; - - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::transformation::OcrDocument; - use serde_json::{Value, json}; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); - } - - #[tokio::test] - async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - "http://127.0.0.1:1", - options.clone(), - ); - let rejected = perform_ocr(request).await.is_err(); - assert!(rejected, "accepted {options}"); - } - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let request_count = seen.clone(); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed - .lock() - .unwrap() - .push((request_count.lock().unwrap().len(), raw.body.clone())); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *responses_received.lock().unwrap(), - [ - (1, r#"{"submitted":true}"#.to_string()), - (2, r#"{"status":"succeeded"}"#.to_string()), - ] - ); - } -} diff --git a/litellm-rust/crates/core/tests/cohere_ocr.rs b/litellm-rust/crates/core/tests/cohere_ocr.rs deleted file mode 100644 index 12824f58b1d..00000000000 --- a/litellm-rust/crates/core/tests/cohere_ocr.rs +++ /dev/null @@ -1,136 +0,0 @@ -mod transformation { - use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, - }, - cohere::ocr::transformation::*, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - #[tokio::test] - async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({ - "output_format":"markdown", "timeout":30, - "extra_body":{ - "output_format": {"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - } - }), - ); - let request = request.with_document( - serde_json::from_value(json!({ - "type":"image_url","image_url":"https://example.com/original.png" - })) - .unwrap(), - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model":"parse", "output_format":{"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - }) - ); - } - - #[tokio::test] - async fn explicit_null_options_use_defaults_before_http() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({"output_format":null,"req_format":null}), - ); - let request = request.with_document( - serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png"}), - ) - .unwrap(), - ); - assert_eq!( - request.response_format().unwrap(), - OcrResponseFormat::Litellm - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!(body["output_format"], "markdown"); - assert!(body.get("req_format").is_none()); - } - - #[rstest] - #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] - #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] - #[tokio::test] - async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( - #[case] model: &str, - #[case] request_line: &str, - ) { - use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = crate::ocr::test_support::wire_request(model, &base, json!({})) - .with_document( - serde_json::from_value::( - json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), - ) - .unwrap() - .into(), - ); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with(request_line), "{}", requests[0]); - assert_eq!( - header(&requests[0], "authorization"), - Some("Bearer test-key") - ); - } - - #[rstest] - #[tokio::test] - async fn route_rejects_non_image_document_without_a_request( - #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, - ) { - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - - let error = perform_ocr(crate::ocr::test_support::wire_request( - model, - &base, - json!({}), - )) - .await - .unwrap_err(); - server.abort(); - - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs deleted file mode 100644 index 96e7451769d..00000000000 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ /dev/null @@ -1,133 +0,0 @@ -use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, - vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, - normalize_response as transform_ocr_response, - }, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -fn document() -> OcrDocument { - serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() -} - -#[rstest] -#[case("stream", json!(true))] -#[case("temperature", json!(0.1))] -#[case("max_tokens", json!(1024))] -#[case("top_p", json!(0.9))] -#[case("n", json!(2))] -#[case("stop", json!("done"))] -#[case("stop", json!(["done", "stop"]))] -fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); -} - -#[rstest] -#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] -#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] -fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); -} - -#[rstest] -#[case(json!("# hello"), "# hello")] -#[case(json!("{broken"), "{broken")] -#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] -#[case(json!({"pages":[]}), "")] -#[case(json!("[]"), "[]")] -#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] -#[case(json!({"pages":[{"markdown":"object"}]}), "object")] -fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) { - let structured = content - .as_object() - .is_some_and(|object| object.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if structured { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } -} - -#[test] -fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = transform_ocr_response("requested", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["index"], 2); - assert_eq!(result["pages"][0]["images"][0]["id"], "one"); - assert_eq!(result["model"], "provider-model"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - assert_eq!(result["document_annotation"]["language"], "en"); - assert_eq!(result["future"], "kept"); -} - -#[test] -fn response_codec_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } -} diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/tests/messages.rs similarity index 79% rename from litellm-rust/crates/core/src/messages/tests.rs rename to litellm-rust/crates/core/tests/messages.rs index ce48752864a..18af8a7d619 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/tests/messages.rs @@ -1,7 +1,11 @@ use std::{sync::Arc, time::Duration}; use futures_util::future::BoxFuture; -use litellm_http::request::{has_bearer_auth, has_header}; +use litellm_core::messages::{ + Error, messages, + route::{LocalMessagesHost, MessagesCall, messages_machine}, + types::{MessagesRequest, MessagesShaping}, +}; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::{Map, Value, json}; use tokio::{ @@ -9,14 +13,6 @@ use tokio::{ net::{TcpListener, TcpStream}, }; -use super::{ - Error, - common_utils::{messages_provider_config, string_headers, truncate_error_body}, - messages, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, -}; -use crate::messages::types::{MessagesRequest, MessagesShaping}; - struct RecordingSecrets { values: Vec<(&'static str, String)>, fails: bool, @@ -73,55 +69,6 @@ fn secrets_call() -> MessagesCall { } } -#[tokio::test] -async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - let secrets = Arc::new(RecordingSecrets::new( - vec![ - ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), - ("ANTHROPIC_BASE_URL", format!("http://{addr}")), - ], - false, - )); - - let output = litellm_host::run::run( - messages_machine(secrets.clone()), - &LocalMessagesHost::new(secrets_call()), - ) - .await - .expect("messages request succeeds"); - - assert!(matches!(output, MessagesOutput::Message(_))); - let request = server.await.expect("server task completes"); - assert!( - request - .to_ascii_lowercase() - .contains("x-api-key: sk-from-manager"), - "{request}" - ); - let requested = secrets.requested.lock().unwrap().clone(); - assert_eq!( - requested, - messages_provider_config("anthropic") - .unwrap() - .secret_names() - .iter() - .map(ToString::to_string) - .collect::>() - ); -} - #[tokio::test] async fn route_surfaces_a_secret_manager_failure_before_the_call() { let Err(error) = litellm_host::run::run( @@ -179,75 +126,6 @@ fn write_response(body: &str) -> String { ) } -#[test] -fn provider_config_resolves_anthropic_and_azure_ai() { - assert!(messages_provider_config("anthropic").is_some()); - assert!(messages_provider_config("azure_ai").is_some()); - assert!(messages_provider_config("openai").is_none()); -} - -#[test] -fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(400); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); -} - -#[test] -fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn has_header_is_case_insensitive() { - let headers = vec![("X-Api-Key".to_string(), "secret".to_string())]; - assert!(has_header(&headers, "x-api-key")); - assert!(!has_header(&headers, "authorization")); -} - -#[test] -fn has_bearer_auth_requires_a_nonempty_bearer_token() { - assert!(has_bearer_auth(&[( - "Authorization".to_string(), - "Bearer tok".to_string() - )])); - assert!(has_bearer_auth(&[( - "authorization".to_string(), - "bearer tok".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Bearer ".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - String::new() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Basic abc".to_string() - )])); - assert!(!has_bearer_auth(&[( - "x-api-key".to_string(), - "sk".to_string() - )])); -} - #[tokio::test] async fn messages_round_trip_builds_azure_request_and_passes_response_through() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs deleted file mode 100644 index 26cd2153e8d..00000000000 --- a/litellm-rust/crates/core/tests/ocr.rs +++ /dev/null @@ -1,1033 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_auth_gcp::VertexAuth; -use litellm_host::{ - event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp, HostResult}, - machine::{HostFailure, Machine, MachineStep}, -}; -use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, -}; -use litellm_llms::base_llm::ocr::{ - error::Error as OcrError, - handler::OcrClient, - settings::OcrSettings, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, - }, -}; -use litellm_secrets::source::SecretSource; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; - -struct RecordingSecretSource { - names: Arc>>, - values: &'static [(&'static str, &'static str)], - api_base: String, -} - -impl SecretSource for RecordingSecretSource { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - self.names.lock().unwrap().push(name.to_owned()); - Box::pin(async move { - Ok(match name { - "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), - _ => self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), - } - .map(litellm_secrets::SecretValue::new)) - }) - } -} - -#[rstest] -#[case::mistral("mistral/model", json!({}))] -#[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] -#[tokio::test] -async fn ocr_contract_upstream_error_preserves_status_body_and_headers( - #[case] model: &str, - #[case] options: Value, -) { - let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); - let expected_body = serde_json::to_string(&payload).unwrap(); - let (base, seen, server) = mock_server(vec![MockResponse { - status: 422, - headers: vec![ - ("Retry-After", "17".into()), - ("X-Request-ID", "request-123".into()), - ("X-Future-Header", "retained".into()), - ], - body: payload, - }]) - .await; - let error = perform_ocr(wire_request(model, &base, options)) - .await - .unwrap_err(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - let OcrError::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; - assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) - ); - } - assert_eq!( - body.len(), - expected_body.len(), - "provider error body was truncated" - ); - assert_eq!(body, expected_body); -} - -#[test] -fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: serde_json::Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); -} - -#[tokio::test] -async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":"ignored" - }) - ); -} - -#[tokio::test] -async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.map(Value::Object), - Some(provider_response) - ); -} - -#[rstest] -#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] -#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] -#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] -#[tokio::test] -async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( - #[case] secrets: &'static [(&'static str, &'static str)], - #[case] expected_key: &str, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: secrets, - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); -} - -#[tokio::test] -async fn mistral_ocr_resolves_provider_secrets_before_transformation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: &[("MISTRAL_API_KEY", "source-key")], - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); -} - -#[tokio::test] -async fn ocr_client_uses_the_injected_http_pool_configuration() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let settings = HttpSettings { - user_agent: Some("host-owned/1".into()), - ..HttpSettings::default() - }; - let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), - ) - .unwrap(); - crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); -} - -fn event_name(event: &CallEvent) -> &'static str { - match event { - CallEvent::Started { .. } => "started", - CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", - CallEvent::Succeeded { .. } => "success", - CallEvent::Failed { .. } => "failure", - } -} - -fn recording_host( - request: crate::ocr::types::LiteLLMOcrRequest, - events: Arc>>, - block: bool, -) -> LocalOcrHost { - let before_send_events = events.clone(); - LocalOcrHost::new(request) - .with_before_send(move |wire, _| { - before_send_events.lock().unwrap().push("before_send"); - if block { - return Err(OcrError::InvalidRequest("blocked".into())); - } - Ok(wire) - }) - .with_observer(move |event| events.lock().unwrap().push(event_name(event))) -} - -#[tokio::test] -async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_before_send( - |mut wire, _| { - wire.headers - .push(("x-core-callback".into(), "edited".into())); - Ok(wire) - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); -} - -#[tokio::test] -async fn before_send_context_names_the_route_and_its_secrets() { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host = LocalOcrHost::new(wire_request( - "mistral/model", - &base, - json!({"pages": [0], "req_format": "native"}), - )) - .with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let (wire, context) = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.custom_llm_provider, "mistral"); - assert_eq!(context.model, "model"); - assert_eq!(wire.body["pages"], json!([0])); - assert!(context.secret_fields.is_empty()); - assert_eq!(context.optional_params["req_format"], "native"); - - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let request = wire_request( - "azure_ai/model", - &base, - json!({"client_secret": "shh", "tenant_id": "t"}), - ); - let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { - bytes: b"abc".as_slice().into(), - file_name: None, - mime_type: Some("application/pdf".into()), - }); - let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some(context.clone()); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let context = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.secret_fields, ["client_secret"]); -} - -#[tokio::test] -async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "response", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", "http://127.0.0.1:1", json!({})), - events.clone(), - true, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); -} - -#[tokio::test] -async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - assert!(perform_ocr_with(host).await.is_err()); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -/// Drives the machine by hand, answering every op through `host` except `before_send`, -/// which `intercept` answers so a test can fail or cancel exactly there. -async fn drive_until( - client: OcrClient, - host: &LocalOcrHost, - mut intercept: impl FnMut(WireRequest) -> Result>, -) -> ( - Result, - Vec<&'static str>, - crate::ocr::route::OcrMachine, -) { - let mut machine = ocr_machine(client); - let mut result = None; - let mut ops = Vec::new(); - let outcome = loop { - let op = match machine.resume(result.take()).await { - Ok(MachineStep::Host(op)) => op, - Ok(MachineStep::Complete(response)) => break Ok(response), - Err(error) => break Err(error), - }; - let answer = match op { - HostOp::Route(op) => { - ops.push(match op { - OcrOp::ProjectRequest => "ProjectRequest", - OcrOp::ReadDocument => "ReadDocument", - OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", - }); - host.route(op) - .await - .map(HostResult::Route) - .map_err(HostFailure::Error) - } - HostOp::BeforeSend { wire, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) - } - HostOp::Emit(event) => { - let event = CallEvent::Machine(event); - ops.push(event_name(&event)); - host.emit(&event) - .await - .map(|()| HostResult::Emitted) - .map_err(HostFailure::Error) - } - }; - match answer { - Ok(answer) => result = Some(answer), - Err(failure) => break machine.interrupt(failure).await, - } - }; - (outcome, ops, machine) -} - -#[tokio::test] -async fn failed_before_send_does_not_replay_or_reach_transport() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Error(OcrError::InvalidRequest( - "before_send failed".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") - ); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(None).await.is_err()); -} - -#[tokio::test] -async fn invalid_provider_response_emits_response_received_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed.lock().unwrap().push(raw.body.clone()); - } - }, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - server.await.unwrap(); - assert!(matches!(error, OcrError::ResponseField { .. })); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - *responses_received.lock().unwrap(), - [r#"{"pages":"invalid"}"#] - ); -} - -#[tokio::test] -async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; - server.await.unwrap(); - assert_eq!(outcome.unwrap().pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); - assert!(matches!( - machine.resume(None).await, - Err(OcrError::InvalidRequest(_)) - )); -} - -async fn drive_native_file_call( - request: crate::ocr::types::LiteLLMOcrRequest, - content: Result, -) -> (Result, usize) { - let reads = Arc::new(Mutex::new(0)); - let counted = reads.clone(); - let content = Mutex::new(Some(content)); - let host = LocalOcrHost::new(request).with_reader(move || { - *counted.lock().unwrap() += 1; - content.lock().unwrap().take().unwrap() - }); - let outcome = perform_ocr_with(host).await; - let reads = *reads.lock().unwrap(); - (outcome, reads) -} - -#[tokio::test] -async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"file"}] - }))]) - .await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::HostReader { - mime_type: Some("application/pdf".into()), - }, - ); - let (response, reads) = drive_native_file_call( - request, - Ok(crate::ocr::types::OcrFileContent { - bytes: b"abc".as_slice().into(), - file_name: Some("scan.png".into()), - }), - ) - .await; - server.await.unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "file"); - assert_eq!(reads, 1); - assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); -} - -#[tokio::test] -async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let failure = OcrError::InvalidRequest("reader exploded".into()); - let (response, reads) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Err(failure.clone()), - ) - .await; - assert!( - matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") - ); - assert_eq!(reads, 1); - - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Ok(crate::ocr::types::OcrFileContent { - bytes: Default::default(), - file_name: None, - }), - ) - .await; - assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn path_documents_are_read_by_core_without_a_host_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"path"}] - }))]) - .await; - let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); - std::fs::create_dir_all(&dir).unwrap(); - let path = dir.join("scan.png"); - std::fs::write(&path, b"abc").unwrap(); - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let (response, reads) = - drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; - server.await.unwrap(); - std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(reads, 0); - assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); - - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }), - Err(OcrError::InvalidRequest("unused".into())), - ) - .await; - assert!(matches!( - response.unwrap_err(), - OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound - )); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Cancelled(OcrError::InvalidRequest( - "cancelled".into(), - ))) - }) - .await; - assert!(matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled")); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); -} - -#[tokio::test] -async fn missing_host_result_preserves_pending_operation() { - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let mut machine = ocr_machine(ocr_client()); - assert!(matches!( - machine.resume(None).await.unwrap(), - MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) - )); - assert!(machine.resume(None).await.is_err()); - assert!(matches!( - machine - .resume(Some(HostResult::Route(OcrOpResult::Request { - request: Box::new(request), - caller_token: false, - }))) - .await - .unwrap(), - MachineStep::Host(HostOp::BeforeSend { .. }) - )); -} - -async fn read_bounded_response(response: Vec, limit: usize) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") -} - -#[tokio::test] -async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use litellm_llms::base_llm::ocr::error::Error; - - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(Error::TooLarge { limit: 8 }) - )); - } -} - -#[rstest] -#[case::declared("Content-Length: 1000000")] -#[case::chunked("Transfer-Encoding: chunked")] -#[tokio::test] -async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( - #[case] headers: &str, -) { - let prefix = "x".repeat(4096); - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), prefix.len()) - .await - .unwrap_err(); - match error { - OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!(body, prefix); - } - error => panic!("unexpected error: {error}"), - } -} - -#[test] -fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.transport.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } -} - -#[derive(Debug)] -struct PendingToken { - entered: Arc, - dropped: Arc, -} - -struct TokenFutureDrop(Arc); - -impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } -} - -impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } -} - -#[tokio::test] -async fn interrupt_drops_provider_captures_before_returning() { - use std::sync::atomic::{AtomicBool, Ordering}; - - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = crate::ocr::types::LiteLLMOcrRequest { - transport: OcrTransportConfig { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.transport - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let host = LocalOcrHost::new(request); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { - HostResult::BeforeSend(wire) - } - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("pending provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = OcrError::InvalidRequest("cancelled".into()); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); - assert!( - dropped.load(Ordering::SeqCst), - "interrupt returned while provider captures were still alive" - ); - assert!( - matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); -} - -struct CallerTokenHost { - request: Mutex>, - trace: Mutex>, -} - -impl Host for CallerTokenHost { - async fn route(&self, op: OcrOp) -> Result { - match op { - OcrOp::ProjectRequest => { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrOpResult::Request { - request: Box::new(self.request.lock().unwrap().take().unwrap()), - caller_token: true, - }) - } - OcrOp::AcquireAzureAdToken => { - self.trace.lock().unwrap().push("token".into()); - Ok(OcrOpResult::AzureAdToken( - litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( - "caller-token", - )), - )) - } - OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), - } - } - - async fn before_send( - &self, - wire: WireRequest, - _: &litellm_host::event::RequestContext, - ) -> Result { - let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); - let authorization = wire - .headers - .iter() - .find(|(name, _)| is_authorization(name)) - .map(|(_, value)| value.clone()) - .unwrap_or_default(); - self.trace - .lock() - .unwrap() - .push(format!("before_send:{authorization}")); - let headers = wire - .headers - .into_iter() - .map(|(name, value)| match is_authorization(&name) { - true => (name, "Bearer edited".to_string()), - false => (name, value), - }) - .collect(); - Ok(WireRequest { headers, ..wire }) - } -} - -#[tokio::test] -async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request("azure_ai/model", &base, json!({})); - request.credentials.api_key = None; - let host = CallerTokenHost { - request: Mutex::new(Some(request)), - trace: Mutex::new(Vec::new()), - }; - - litellm_host::run::run(ocr_machine(ocr_client()), &host) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!( - *host.trace.lock().unwrap(), - ["project", "token", "before_send:Bearer caller-token"] - ); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer edited\r\n") - ); -} - -#[tokio::test] -async fn interrupting_an_in_flight_provider_request_closes_its_connection() { - use tokio::io::AsyncReadExt; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let received = Arc::new(tokio::sync::Notify::new()); - let server_received = received.clone(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = Vec::new(); - let mut buffer = [0u8; 4096]; - while !request.windows(4).any(|window| window == b"\r\n\r\n") { - let read = socket.read(&mut buffer).await.unwrap(); - request.extend_from_slice(&buffer[..read]); - } - server_received.notify_one(); - loop { - if socket.read(&mut buffer).await.unwrap() == 0 { - break; - } - } - }); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = received.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("the stalled provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - - let cancelled = OcrError::InvalidRequest("cancelled".into()); - assert!( - machine - .interrupt(HostFailure::Cancelled(cancelled)) - .await - .is_err() - ); - tokio::time::timeout(std::time::Duration::from_secs(1), server) - .await - .expect("the provider connection stayed open after the interrupt") - .unwrap(); -} diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs deleted file mode 100644 index 855548dc6bf..00000000000 --- a/litellm-rust/crates/core/tests/ocr/document.rs +++ /dev/null @@ -1,152 +0,0 @@ -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send(self, wire: WireRequest) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - result: Result<(), Error>, - provider_body: Option, -} - -async fn send(route: Route, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document_type = route.document_type(); - let document = - json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = - LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - Sent { - result, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::Detached, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[rstest] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::ReplacesDocument, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs deleted file mode 100644 index 974fa3d6655..00000000000 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient}, - transformation::LiteLLMOcrResponse, -}; -use serde_json::{Value, json}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, -}; - -use crate::ocr::{ - route::{LocalOcrHost, ocr_machine}, - types::LiteLLMOcrRequest, - wire::{OcrWireRequest, decode_request}, -}; - -/// Stands in for a host with no hooks registered: the wire request goes out unchanged -/// and response events go nowhere. -pub(crate) struct NoHooks; - -impl CallHooks for NoHooks { - fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(wire) }) - } - - fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(async { Ok(()) }) - } -} - -pub(crate) fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) -} - -pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result { - crate::ocr::client::perform(&ocr_client(), request).await -} - -pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { - litellm_host::run::run(ocr_machine(ocr_client()), &host).await -} - -pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - base, - json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - options, - ) -} - -pub(crate) fn wire_request_with_document( - model: &str, - base: &str, - document: Value, - options: Value, -) -> LiteLLMOcrRequest { - decode_request(OcrWireRequest { - model: model.into(), - document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), - api_base: Some(base.into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap() -} - -pub(crate) fn resolved_request( - request: LiteLLMOcrRequest, -) -> crate::ocr::types::ResolvedOcrRequest { - request - .map_document(crate::ocr::document::prepare_document) - .unwrap() -} - -pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { - let request = resolved_request(request); - let document = request.document.clone().with_source(source.into()); - request.with_document(document.into()) -} - -pub(crate) fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; - -/// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. -pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - loop { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut buffer = [0u8; 4096]; - let _ = socket.read(&mut buffer).await.unwrap(); - let head = format!( - "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - SERVED_DOCUMENT.len() - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(SERVED_DOCUMENT).await.unwrap(); - } - }); - (base, task) -} - -pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, -} - -impl MockResponse { - pub fn json(body: Value) -> Self { - Self { - status: 200, - headers: vec![], - body, - } - } -} - -pub(crate) async fn mock_server( - responses: Vec, -) -> (String, Arc>>, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let requests = Arc::new(Mutex::new(Vec::new())); - let seen = requests.clone(); - let server_base = base.clone(); - let task = tokio::spawn(async move { - for response in responses { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0u8; 4096]; - let header_end = loop { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { - break index + 4; - } - }; - let length = String::from_utf8_lossy(&bytes[..header_end]) - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap_or(0); - while bytes.len() < header_end + length { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - } - seen.lock() - .unwrap() - .push(String::from_utf8_lossy(&bytes).into_owned()); - let body = serde_json::to_vec(&response.body).unwrap(); - let headers = response - .headers - .into_iter() - .map(|(name, value)| { - format!("{name}: {}\r\n", value.replace("{base}", &server_base)) - }) - .collect::(); - let head = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", - response.status, - body.len(), - headers - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(&body).await.unwrap(); - } - }); - (base, requests, task) -} - -pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { - request - .lines() - .take_while(|line| !line.is_empty()) - .find_map(|line| { - let (key, value) = line.split_once(':')?; - key.eq_ignore_ascii_case(name).then(|| value.trim()) - }) -} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs deleted file mode 100644 index 83e7754122b..00000000000 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ /dev/null @@ -1,584 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; -use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[rstest] -#[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[tokio::test] -async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = super::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); -} - -#[rstest] -#[case("parse-v3")] -#[case("parse-legacy")] -#[tokio::test] -async fn data_uri_upload_preserves_multipart_headers( - #[case] model: &str, - #[values("application/pdf", "image/png")] mime_type: &str, -) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let document = if mime_type.starts_with("image/") { - json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) - } else { - json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) - }; - let mut request = crate::ocr::types::LiteLLMOcrRequest { - document: serde_json::from_value::(document) - .unwrap() - .into(), - ..wire_request(&format!("reducto/{model}"), &base, json!({})) - }; - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; - assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); - assert!(multipart.contains("\r\n\r\nabc\r\n--")); - assert!(requests[1].starts_with("POST /parse ")); - let source_field = if model == "parse-legacy" { - "document_url" - } else { - "input" - }; - assert_eq!( - request_body(&requests[1]), - json!({source_field:"reducto://uploaded.pdf"}) - ); - for request in requests.iter() { - assert!( - request - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - } -} - -#[tokio::test] -async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[rstest] -#[case(json!({"file_id":""}))] -#[case(json!({}))] -#[case(json!({"file_id":null}))] -#[tokio::test] -async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[rstest] -#[case("https://example.com/a.pdf", Error::ReductoSource)] -#[case("reducto://", Error::RequestField { path: "document file id".into() })] -#[case("data:application/pdf;base64", Error::InvalidDataUri)] -#[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] -#[tokio::test] -async fn rejects_invalid_document_sources_before_network( - #[case] source: &str, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - source, - ); - let result = perform_ocr(request).await; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid source: {source}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[test] -fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use litellm_llms::reducto::ocr::transformation::{ - ReductoResponse, normalize_response as transform_ocr_response, - }; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = transform_ocr_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = transform_ocr_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); -} - -#[tokio::test] -async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); -} - -#[tokio::test] -async fn native_format_retains_the_provider_response() { - let raw = json!({ - "result":{"chunks":[{"content":"native OCR response"}]}, - "usage":{"num_pages":1} - }); - let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages[0].markdown, "native OCR response"); - assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); -} - -#[tokio::test] -async fn unknown_model_reaches_parse_and_keeps_its_name() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[{"content":"future model response"}]} - }))]) - .await; - let request = super::test_support::with_source( - wire_request("reducto/future-parse-model", &base, json!({})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.model, "future-parse-model"); - assert_eq!(response.pages[0].markdown, "future model response"); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!( - request_body(&requests[0]), - json!({"input":"reducto://ready.pdf"}) - ); -} - -#[tokio::test] -async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_before_send(|wire, _| { - assert_eq!( - wire.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(WireRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); -} - -mod transformation { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, - reducto::ocr::transformation::*, - }; - use rstest::rstest; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn v3_options_preserve_explicit_null() { - let overrides = - serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) - .unwrap(); - let params = ReductoParseV3Config - .map_ocr_params(&overrides, "parse-v3") - .unwrap(); - let client = crate::ocr::test_support::ocr_client(); - let connection = OcrConnection::default(); - let document = serde_json::from_value( - json!({"type":"document_url","document_url":"reducto://ready.pdf"}), - ) - .unwrap(); - let body = ReductoParseV3Config - .async_transform_ocr_request( - "parse-v3", - document, - ¶ms, - &[], - OcrRequestContext { - client: &client, - connection: &connection, - }, - ) - .await - .unwrap(); - assert_eq!( - serde_json::to_value(body).unwrap(), - json!({ - "input":"reducto://ready.pdf", "formatting":null, "settings":{} - }) - ); - let absent = ReductoParseV3Config - .map_ocr_params( - &litellm_core_utils::call_arguments::CallArguments::default(), - "parse-v3", - ) - .unwrap(); - assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case("https://example.com/a.pdf")] - #[case("reducto://")] - #[case("data:application/pdf;base64")] - #[case("data:application/pdf;base64,INVALID!")] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), - source, - ); - assert!(perform_ocr(request).await.is_err()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[rstest] - #[case("reducto/parse-v3")] - #[case("reducto/parse-legacy")] - #[tokio::test] - async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())]; - let host = LocalOcrHost::new(request).with_before_send(|wire, _| { - Ok(WireRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs deleted file mode 100644 index 2e8d69f5f64..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ /dev/null @@ -1,143 +0,0 @@ -use litellm_auth::InputSource; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = super::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); -} - -#[test] -fn host_registration_selects_deepseek_without_affecting_mistral() { - assert!(crate::ocr::arguments::is_supported_request( - "deepseek-ocr-maas", - Some("vertex_ai") - )); - assert!(crate::ocr::arguments::is_supported_request( - "mistral-ocr-maas", - Some("vertex_ai") - )); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/deepseek-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -mod deepseek_transformation { - use serde_json::json; - - use super::*; - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs deleted file mode 100644 index 035f3fe944d..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_auth::InputSource; -use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, ocr_client, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); -} - -#[tokio::test] -async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let client = ocr_client().with_settings(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - crate::ocr::client::perform( - &client, - wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].starts_with( - "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); -} - -#[tokio::test] -async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); -} - -#[tokio::test] -async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -#[tokio::test] -async fn adapters_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "ignored" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - for http in [&direct_http, &vertex_http] { - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "ignored" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); -} - -mod transformation { - - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::test_support::wire_request; - - #[rstest] - #[case::mistral(false)] - #[case::vertex(true)] - #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization( - #[case] use_vertex: bool, - ) { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - let http = if use_vertex { - &vertex_http - } else { - &direct_http - }; - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "preserved" - }) - ); - let payload = serde_json::to_vec( - &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), - ) - .unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &payload, Default::default()) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &payload, Default::default()) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } -} diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index 874a0f3abf9..fcf296793a5 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -217,13 +217,25 @@ mod tests { "Authorization".to_string(), "Bearer abc".to_string() )])); + assert!(has_bearer_auth(&[( + "authorization".to_string(), + "bearer abc".to_string() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Bearer ".to_string() )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + String::new() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Basic abc".to_string() )])); + assert!(!has_bearer_auth(&[( + "x-api-key".to_string(), + "abc".to_string() + )])); } } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 6fc4f00b981..fd86c5ca25a 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -218,7 +218,3 @@ fn anthropic_body( ); Value::Object(body) } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 09c456f1d0a..b5db88d7dc4 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -302,7 +302,3 @@ fn has_blank_text(message: &ChatMessage) -> bool { }), } } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs similarity index 96% rename from litellm-rust/crates/llms/src/anthropic/chat/tests.rs rename to litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs index 3777347d240..ed22a1d141d 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") @@ -205,7 +209,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "image_url", "image_url": {"url": "https://x/y.png"}} ]}]), json!({}) @@ -214,7 +219,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}} ]}]), json!({}) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs similarity index 98% rename from litellm-rust/crates/llms/src/bedrock/chat/tests.rs rename to litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs index d7ecde47c6b..4127bcfa19d 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, + bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") diff --git a/litellm-rust/crates/types/src/utils.rs b/litellm-rust/crates/types/src/utils.rs index 5ca56ec9e49..af0ba2c01c9 100644 --- a/litellm-rust/crates/types/src/utils.rs +++ b/litellm-rust/crates/types/src/utils.rs @@ -56,7 +56,7 @@ pub struct ChatCompletionsChoice { /// /// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the /// `ModelResponse` it already created, and echoing the provider's own id here -/// would change it. Pinned by `response_carries_no_id` in `tests.rs`. +/// would change it. Pinned by `response_carries_no_id` in the Anthropic chat transformation tests. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionsResponse { pub created: u64, From e64e635185b25cb9ee649dba70a372a869270264 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:35:06 -0700 Subject: [PATCH 008/154] fix(cost-map): add video and reasoning output prices to vertex gemini-omni-1.1-flash (#43036) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dad89fc58a8..90a268aa3e2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68612,7 +68612,9 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai", "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/gemini-omni-1.1-flash-preview": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dad89fc58a8..90a268aa3e2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68612,7 +68612,9 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai", "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/gemini-omni-1.1-flash-preview": { From 77eccaca78d25820c54c6b2703711efe3f95224f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:40:36 -0500 Subject: [PATCH 009/154] feat(proxy): server-side Team Usage export beyond the top-N key cap (#42996) * feat(proxy): add uncapped server-side team usage export route GET /team/daily/activity/export answers the same scoping as /team/daily/activity/aggregated with one unbounded rollup query, so keys past USAGE_TOP_API_KEYS_LIMIT are included. Supports daily, daily_with_keys, daily_with_users and daily_with_models export types as CSV (default) or JSON. The PTU flat-cost sentinel stays in the plain daily rollup and is excluded from the keyed and per-model exports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): export team usage server-side when the key list was truncated When the aggregated spend response reports api_key truncation, EntityUsage passes a serverExport into the export modal that downloads CSV or JSON from GET /team/daily/activity/export instead of building the file from the truncated on-screen data. apiClient gains a responseType option so the download can arrive as a Blob, and truncation no longer blocks the export button Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover team usage export types, sentinel handling and scope Unit tests pin the uncapped key rollup past USAGE_TOP_API_KEYS_LIMIT, PTU sentinel inclusion in the daily rollup and exclusion elsewhere, the per-user fold, and the CSV column layout. Integration tests exercise the route against a live proxy, including member scope denial Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): tidy team usage export route Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): use membership test for export type branch (PLR1714) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): format exportBlockedReason test with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): satisfy type-discipline gate in team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): pass export rows as a sequence to the response model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-behavior): cover team usage export in the daily activity scope matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(proxy): carry PTU flat cost and escape formulas in team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): keep the truncation export block on surfaces without a server export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): drop redundant comments in team export call and modal test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit cells for team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten team usage export audit cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type the export params tuple and fold user keys in one pass Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): bring entity usage export helpers under eslint budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): prettier-format UsagePageView after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 4 + .../common_daily_activity.py | 293 +++++++++- .../management_endpoints/team_endpoints.py | 183 +++++- .../management_endpoints/team_endpoints.py | 45 ++ .../endpointaudit/coverage_allowlist.txt | 1 + .../spend/test_team_daily_activity_export.py | 522 ++++++++++++++++++ .../management/test_team_daily_activity.py | 12 +- .../test_common_daily_activity.py | 305 ++++++++++ .../test_team_endpoints.py | 127 +++++ .../components/EntityUsage/EntityUsage.tsx | 19 +- .../_components/components/UsagePageView.tsx | 11 +- .../EntityUsageExportModal.test.tsx | 31 +- .../EntityUsageExportModal.tsx | 8 +- .../EntityUsageExport/UsageExportHeader.tsx | 5 +- .../exportBlockedReason.test.ts | 4 + .../EntityUsageExport/exportBlockedReason.ts | 4 +- .../src/components/EntityUsageExport/types.ts | 3 + .../EntityUsageExport/utils.test.ts | 37 ++ .../src/components/EntityUsageExport/utils.ts | 71 ++- .../src/components/networking.tsx | 31 ++ ui/litellm-dashboard/src/lib/http/client.ts | 10 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 147 ++++- 22 files changed, 1812 insertions(+), 61 deletions(-) create mode 100644 tests/integration/spend/test_team_daily_activity_export.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 73e3e0e6ee0..b6de36f8423 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -307,6 +307,7 @@ class KeyManagementRoutes(str, enum.Enum): # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated" + TEAM_DAILY_ACTIVITY_EXPORT = "/team/daily/activity/export" TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search" # team spend-log viewing @@ -719,6 +720,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only @@ -890,6 +892,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/team/spend/by_user", "/team/{team_id}/members/me", @@ -990,6 +993,7 @@ class LiteLLMRoutes(enum.Enum): "/user/daily/activity", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/tag/daily/activity", "/tag/list", diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index bf8e7bc15fb..1d797206ead 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,4 +1,6 @@ import asyncio +import dataclasses +import itertools from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone @@ -35,6 +37,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, SpendMetrics, ) +from litellm.types.proxy.management_endpoints.team_endpoints import ( + TeamDailyActivityExportRow, + TeamDailyActivityExportType, +) if TYPE_CHECKING: from prisma.models import ( @@ -198,7 +204,7 @@ class _AggregatedQueryKwargs(TypedDict): include_current_utc_day: ReadOnly[bool] -_SqlQuery = tuple[str, list[str]] +_SqlQuery = tuple[str, Sequence[str]] async def _query_raw_optional( @@ -974,6 +980,291 @@ def _build_entity_rollup_sql_query( return sql_query, sql_params +def _build_export_sql_query( + *, + table_name: str, + entity_id_field: str, + entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + start_date: str, + end_date: str, + api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path + timezone_offset_minutes: int | None, + export_type: TeamDailyActivityExportType, +) -> tuple[str, tuple[str, ...]]: + """One unbounded rollup for the export route, on the aggregated path's WHERE clause. + + No LIMIT anywhere: the export exists so a caller can reach keys past + USAGE_TOP_API_KEYS_LIMIT. PTU sentinel rows stay in `daily` so per-team + totals match breakdown.entities, and are excluded from the key, user and + model exports where the flat-cost row has no meaning. + """ + pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name) + if pg_table is None: + raise ValueError(f"Unknown table name: {table_name}") + + adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) + where_clause, where_params = _build_aggregated_where_clause( + entity_id_field=entity_id_field, + entity_id=entity_id, + adjusted_start=adjusted_start, + adjusted_end=adjusted_end, + model=None, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + ) + + keyed: Final = export_type in ("daily_with_keys", "daily_with_users") + by_model: Final = export_type == "daily_with_models" + group_extras: Final = tuple(field for field in ("api_key" if keyed else "", "model" if by_model else "") if field) + group_by: Final = f'date, "{entity_id_field}"' + "".join(f", {field}" for field in group_extras) + sentinel_clause: Final = f" AND api_key <> ${len(where_params) + 1}" if (keyed or by_model) else "" + sentinel_params: Final = (PTU_SENTINEL_API_KEY,) if (keyed or by_model) else () + + sql_query: Final = f""" + SELECT + date, + "{entity_id_field}" AS entity_id, + {"api_key" if keyed else "NULL::text AS api_key"}, + {"model" if by_model else "NULL::text AS model"},{_rollup_metric_select(table_name)} + FROM "{pg_table}" + WHERE {where_clause}{sentinel_clause} + GROUP BY {group_by} + ORDER BY {group_by} + """ + + return sql_query, (*where_params, *sentinel_params) + + +class _ExportRow(_RollupMetricsRow): + entity_id: str | None + model: str | None + + +def _export_team_alias(entity_metadata_field: Mapping[str, dict[str, object]] | None, entity_id: str) -> str | None: + alias: Final = _entity_metadata(entity_metadata_field, entity_id).get("team_alias") + return alias if isinstance(alias, str) else None + + +@dataclasses.dataclass(frozen=True, slots=True) +class _ExportMetrics: + spend: float + api_requests: int + successful_requests: int + failed_requests: int + total_tokens: int + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + + @classmethod + def from_record(cls, record: _RollupMetricsRow) -> "_ExportMetrics": + prompt_tokens: Final = record.prompt_tokens or 0 + completion_tokens: Final = record.completion_tokens or 0 + return cls( + spend=record.spend or 0.0, + api_requests=record.api_requests or 0, + successful_requests=record.successful_requests or 0, + failed_requests=record.failed_requests or 0, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cache_read_input_tokens=record.cache_read_input_tokens or 0, + cache_creation_input_tokens=record.cache_creation_input_tokens or 0, + ) + + @classmethod + def zero(cls) -> "_ExportMetrics": + return cls( + spend=0.0, + api_requests=0, + successful_requests=0, + failed_requests=0, + total_tokens=0, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + + def __add__(self, other: "_ExportMetrics") -> "_ExportMetrics": + return _ExportMetrics( + spend=self.spend + other.spend, + api_requests=self.api_requests + other.api_requests, + successful_requests=self.successful_requests + other.successful_requests, + failed_requests=self.failed_requests + other.failed_requests, + total_tokens=self.total_tokens + other.total_tokens, + prompt_tokens=self.prompt_tokens + other.prompt_tokens, + completion_tokens=self.completion_tokens + other.completion_tokens, + cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens, + cache_creation_input_tokens=self.cache_creation_input_tokens + other.cache_creation_input_tokens, + ) + + +def _export_base_row( + record: _ExportRow, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> TeamDailyActivityExportRow: + entity_id: Final = record.entity_id or "Unassigned" + metrics: Final = _ExportMetrics.from_record(record) + return TeamDailyActivityExportRow( + date=record.date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + model=record.model, + spend=metrics.spend, + flat_cost=_reported_flat_cost(record), + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +def _export_key_row( + record: _ExportRow, + entity_metadata_field: Mapping[str, dict[str, object]] | None, + api_key_metadata: Mapping[str, _KeyMetadataDict], +) -> TeamDailyActivityExportRow: + entity_id: Final = record.entity_id or "Unassigned" + metadata: Final = _key_metadata(api_key_metadata, record.api_key or "") + metrics: Final = _ExportMetrics.from_record(record) + return TeamDailyActivityExportRow( + date=record.date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + api_key=record.api_key, + key_alias=metadata.key_alias, + user_id=metadata.user_id, + user_email=metadata.user_email, + spend=metrics.spend, + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +def _fold_export_users( + records: Sequence[_ExportRow], + entity_metadata_field: Mapping[str, dict[str, object]] | None, + api_key_metadata: Mapping[str, _KeyMetadataDict], +) -> tuple[TeamDailyActivityExportRow, ...]: + """Fold (date, team, api_key) rows into (date, team, user) rows.""" + + def bucket_of(record: _ExportRow) -> tuple[str, str, str]: + return ( + record.date, + record.entity_id or "Unassigned", + _key_metadata(api_key_metadata, record.api_key or "").user_id or "Unassigned", + ) + + key_sets: Final = MappingProxyType( + { + bucket: frozenset(record.api_key or "" for record in group) + for bucket, group in itertools.groupby(sorted(records, key=bucket_of), key=bucket_of) + } + ) + sums: Final[dict[tuple[str, str, str], _ExportMetrics]] = {} # mutable-ok: local fold accumulator + emails: Final[dict[tuple[str, str, str], str | None]] = {} # mutable-ok: local fold accumulator + for record in records: + metadata = _key_metadata(api_key_metadata, record.api_key or "") + bucket_key = bucket_of(record) + sums[bucket_key] = sums.get(bucket_key, _ExportMetrics.zero()) + _ExportMetrics.from_record(record) + emails.setdefault(bucket_key, metadata.user_email) + if emails[bucket_key] is None and metadata.user_email is not None: + emails[bucket_key] = metadata.user_email + return tuple( + _export_folded_user_row( + bucket_key, sums[bucket_key], emails[bucket_key], len(key_sets[bucket_key]), entity_metadata_field + ) + for bucket_key in sorted(sums) + ) + + +def _export_folded_user_row( + bucket_key: tuple[str, str, str], + metrics: _ExportMetrics, + user_email: str | None, + keys: int, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> TeamDailyActivityExportRow: + date, entity_id, user_id = bucket_key + return TeamDailyActivityExportRow( + date=date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + user_id=user_id if user_id != "Unassigned" else None, + user_email=user_email, + keys=keys, + spend=metrics.spend, + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +async def get_daily_activity_export_rows( + *, + prisma_client: PrismaClient, + table_name: str, + entity_id_field: str, + entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + entity_metadata_field: Mapping[str, dict[str, object]] | None, + start_date: str, + end_date: str, + api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path + timezone_offset_minutes: int | None, + export_type: TeamDailyActivityExportType, +) -> tuple[TeamDailyActivityExportRow, ...]: + """Every (date, entity[, api_key|model]) rollup row in the range, uncapped.""" + sql_query, sql_params = _build_export_sql_query( + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + start_date=start_date, + end_date=end_date, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + timezone_offset_minutes=timezone_offset_minutes, + export_type=export_type, + ) + raw_rows: Final = await _query_raw_optional(prisma_client, (sql_query, sql_params)) + records: Final = tuple(_ExportRow(**row) for row in (raw_rows or ())) + + if export_type in ("daily", "daily_with_models"): + return await asyncio.to_thread( + lambda: tuple(_export_base_row(record, entity_metadata_field) for record in records) + ) + + api_keys: Final = frozenset(record.api_key for record in records if record.api_key) + api_key_metadata: Final = ( + await get_api_key_metadata(prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records))) + if api_keys + else _EMPTY_KEY_METADATA + ) + if export_type == "daily_with_keys": + return await asyncio.to_thread( + lambda: tuple(_export_key_row(record, entity_metadata_field, api_key_metadata) for record in records) + ) + return await asyncio.to_thread(_fold_export_users, records, entity_metadata_field, api_key_metadata) + + def _aggregate_spend_records_sync( *, records: Sequence[DailySpendRecord], diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 493c83c730a..8a6cd1218ee 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -11,6 +11,8 @@ All /team management endpoints import asyncio import copy +import csv +import io import json import math import traceback @@ -33,7 +35,8 @@ from typing import ( ) import fastapi -from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, Response, status +from fastapi.responses import JSONResponse from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never @@ -126,6 +129,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, + get_daily_activity_export_rows, ) from litellm.proxy.management_endpoints.common_utils import ( _check_disable_global_guardrails_caller_permission, @@ -206,6 +210,11 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, + TeamDailyActivityExportFormat, + TeamDailyActivityExportMetadata, + TeamDailyActivityExportResponse, + TeamDailyActivityExportRow, + TeamDailyActivityExportType, TeamIdSearchFilter, TeamIdSearchMatch, TeamKeyActivitySearchWhere, @@ -6809,6 +6818,178 @@ async def get_team_daily_activity_aggregated( ) +_EXPORT_CSV_METRIC_HEADERS: Final = ( + "Spend ($)", + "Requests", + "Successful Requests", + "Failed Requests", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", +) + + +def _export_csv_headers(export_type: TeamDailyActivityExportType) -> tuple[str, ...]: + base: Final = ("Date", "Team", "Team ID") + if export_type == "daily_with_keys": + return (*base, "Key Alias", "Key ID", "User ID", "User Email", *_EXPORT_CSV_METRIC_HEADERS) + if export_type == "daily_with_users": + return (*base, "User ID", "User Email", "Keys", *_EXPORT_CSV_METRIC_HEADERS) + if export_type == "daily_with_models": + return ( + *base, + "Model", + "Spend ($)", + "Requests", + "Successful", + "Failed", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", + ) + return (*base, *_EXPORT_CSV_METRIC_HEADERS) + + +def _csv_safe(value: str) -> str: + return "'" + value if value[:1] in ("=", "+", "-", "@", "\t", "\r") else value + + +def _export_csv_record(row: TeamDailyActivityExportRow) -> dict[str, object]: + return { # mutable-ok: csv.DictWriter consumes a plain mapping per row + "Date": row.date, + "Team": _csv_safe(row.team_alias) if row.team_alias else "-", + "Team ID": row.team_id, + "Key Alias": _csv_safe(row.key_alias) if row.key_alias else "-", + "Key ID": row.api_key or "-", + "User ID": _csv_safe(row.user_id) if row.user_id else "-", + "User Email": _csv_safe(row.user_email) if row.user_email else "-", + "Keys": row.keys, + "Model": _csv_safe(row.model) if row.model else "-", + "Spend ($)": f"{row.spend:.4f}", + "Flat Cost ($)": f"{row.flat_cost:.4f}", + "Total Cost ($)": f"{row.spend + row.flat_cost:.4f}", + "Requests": row.api_requests, + "Successful Requests": row.successful_requests, + "Failed Requests": row.failed_requests, + "Successful": row.successful_requests, + "Failed": row.failed_requests, + "Total Tokens": row.total_tokens, + "Prompt Tokens": row.prompt_tokens, + "Completion Tokens": row.completion_tokens, + "Cache Read Input Tokens": row.cache_read_input_tokens, + "Cache Creation Input Tokens": row.cache_creation_input_tokens, + } + + +def _team_export_csv(export_type: TeamDailyActivityExportType, rows: Sequence[TeamDailyActivityExportRow]) -> str: + base_headers: Final = _export_csv_headers(export_type) + spend_index: Final = base_headers.index("Spend ($)") + 1 + headers: Final = ( + (*base_headers[:spend_index], "Flat Cost ($)", "Total Cost ($)", *base_headers[spend_index:]) + if sum(row.flat_cost for row in rows) > 0 + else base_headers + ) + buffer: Final = io.StringIO() + writer: Final = csv.DictWriter(buffer, fieldnames=headers, extrasaction="ignore") + writer.writeheader() + writer.writerows(_export_csv_record(row) for row in rows) + return buffer.getvalue() + + +@router.get( + "/team/daily/activity/export", + response_model=TeamDailyActivityExportResponse, + responses={200: {"content": {"text/csv": {}, "application/json": {}}}}, # mutable-ok: OpenAPI content map + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list +) +async def get_team_daily_activity_export( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_date: str | None = None, + end_date: str | None = None, + export_type: TeamDailyActivityExportType = "daily", + format: TeamDailyActivityExportFormat = "csv", + team_id: str | None = None, + exclude_team_ids: str | None = None, + timezone_offset: Annotated[int | None, Query(alias="timezone")] = None, +) -> Response: + """ + Server-side Team Usage export, not subject to USAGE_TOP_API_KEYS_LIMIT. + + Same scoping as /team/daily/activity/aggregated, answered by one unbounded + rollup query, returned as CSV or JSON. For daily_with_keys, + daily_with_users and daily_with_models the PTU sentinel flat-cost rows are + excluded, so metadata totals under those export types cover request spend + only; the plain daily export includes them. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) + + range_error: Final = _aggregated_date_range_error(start_date, end_date) + if range_error is not None or start_date is None or end_date is None: + raise _daily_activity_error(status_code=400, message=range_error or "Please provide start_date and end_date") + + scope: Final = await _resolve_team_daily_activity_scope( + team_ids=team_id, + exclude_team_ids=exclude_team_ids, + api_key=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + rows: Final = await get_daily_activity_export_rows( + prisma_client=prisma_client, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=scope.team_ids, + entity_metadata_field=scope.team_alias_metadata, + start_date=start_date, + end_date=end_date, + api_key=scope.api_key_filter, + exclude_entity_ids=scope.exclude_team_ids, + timezone_offset_minutes=timezone_offset, + export_type=export_type, + ) + + now: Final = datetime.now(timezone.utc) + metadata: Final = TeamDailyActivityExportMetadata( + export_date=now.isoformat(), + export_type=export_type, + start_date=start_date, + end_date=end_date, + team_ids=list(scope.team_ids) if scope.team_ids else None, # mutable-ok: response model field type + total_spend=sum(row.spend for row in rows), + total_flat_cost=sum(row.flat_cost for row in rows), + total_api_requests=sum(row.api_requests for row in rows), + total_successful_requests=sum(row.successful_requests for row in rows), + total_failed_requests=sum(row.failed_requests for row in rows), + total_tokens=sum(row.total_tokens for row in rows), + ) + + if format == "json": + return JSONResponse( + content=TeamDailyActivityExportResponse(metadata=metadata, data=rows).model_dump(mode="json") + ) + return Response( + content=_team_export_csv(export_type, rows), + media_type="text/csv; charset=utf-8", + headers={ # mutable-ok: starlette Response headers is a dict + "Content-Disposition": f'attachment; filename="team_usage_{export_type}_{now.date().isoformat()}.csv"' + }, + ) + + def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere: """Caller scoping lives inside the same Prisma where as the search term so `take` never trims visible matches in favour of keys the caller is not allowed to see.""" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index aac2703e918..9b87d114a81 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -277,3 +277,48 @@ class TeamUserSpendResponse(BaseModel): start_date: str end_date: str results: tuple[TeamUserSpendRow, ...] + + +TeamDailyActivityExportType = Literal["daily", "daily_with_keys", "daily_with_users", "daily_with_models"] +TeamDailyActivityExportFormat = Literal["csv", "json"] + + +class TeamDailyActivityExportRow(BaseModel): + date: str + team_id: str + team_alias: str | None = None + api_key: str | None = None + key_alias: str | None = None + user_id: str | None = None + user_email: str | None = None + keys: int | None = None + model: str | None = None + spend: float + flat_cost: float = 0.0 + api_requests: int + successful_requests: int + failed_requests: int + total_tokens: int + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + + +class TeamDailyActivityExportMetadata(BaseModel): + export_date: str + export_type: TeamDailyActivityExportType + start_date: str + end_date: str + team_ids: list[str] | None + total_spend: float + total_flat_cost: float = 0.0 + total_api_requests: int + total_successful_requests: int + total_failed_requests: int + total_tokens: int + + +class TeamDailyActivityExportResponse(BaseModel): + metadata: TeamDailyActivityExportMetadata + data: list[TeamDailyActivityExportRow] diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index b0f28a9c740..7aacf7ceab9 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated +GET /team/daily/activity/export GET /team/daily/activity/aggregated/search GET /team/spend/by_user GET /team/spend/report diff --git a/tests/integration/spend/test_team_daily_activity_export.py b/tests/integration/spend/test_team_daily_activity_export.py new file mode 100644 index 00000000000..b35d3fe0c8a --- /dev/null +++ b/tests/integration/spend/test_team_daily_activity_export.py @@ -0,0 +1,522 @@ +import csv +import io +import os +import signal +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import group_members, owned_proxy, owned_proxy_process + + +def _export_range() -> dict[str, str]: + today: Final = datetime.now(timezone.utc) + return { + "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), + "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"), + "timezone": "0", + } + + +def _team_with_three_keys( + gateway: Gateway, scenario: Scenario, model: str +) -> tuple[str, tuple[str, ...], tuple[str, ...], dict[str, float]]: + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + return team, keys, digests, spend_by_key + + +def _export_json(gateway: Gateway, **params: str) -> httpx.Response: + return gateway.request("GET", "/team/daily/activity/export", params={**_export_range(), **params}) + + +def test_team_activity_export_returns_every_key_beyond_the_top_n_cap(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={ + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "json", + }, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + rows: Final = tuple(object_value(row) for row in body["data"]) + assert sorted(string_value(row["api_key"]) for row in rows) == sorted(digests), response.text + for row in rows: + assert row["team_id"] == team, response.text + assert float(row["spend"]) == pytest.approx(spend_by_key[string_value(row["api_key"])]), response.text + metadata: Final = object_value(body["metadata"]) + assert ( + metadata["export_type"], + metadata["team_ids"], + metadata["total_api_requests"], + metadata["total_successful_requests"], + metadata["total_failed_requests"], + ) == ("daily_with_keys", [team], 3, 3, 0), response.text + assert float(metadata["total_spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_team_activity_export_csv_downloads_every_key(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={ + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "csv", + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/csv"), response.headers + assert "attachment" in response.headers["content-disposition"], response.headers + records: Final = tuple(csv.DictReader(io.StringIO(response.text))) + assert len(records) == 3, response.text + assert sorted(record["Key ID"] for record in records) == sorted(digests), response.text + assert sorted(record["Team ID"] for record in records) == [team, team, team], response.text + for record in records: + assert record["Spend ($)"] == f"{spend_by_key[record['Key ID']]:.4f}", response.text + + +def test_team_activity_export_denies_a_member_another_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + member: Final = scenario.user(user_role="internal_user", teams=[team_a]) + member_key: Final = scenario.key(user_id=member, team_id=team_a, models=[model]) + reply: Final = gateway.chat(model, key=member_key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_a,)), + lambda values: len(values) == 1, + seconds=70, + ) + denied: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "team_id": team_b, "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert denied.status_code == 404, denied.text + assert f"User does not belong to Team= {team_b}" in denied.text, denied.text + allowed: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "team_id": team_a, "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert allowed.status_code == 200, allowed.text + rows: Final = tuple(object_value(row) for row in object_value(allowed.json())["data"]) + assert len(rows) == 1, allowed.text + assert rows[0]["team_id"] == team_a, allowed.text + assert float(rows[0]["spend"]) == pytest.approx(float(daily[0]["spend"])), allowed.text + + +def test_export_daily_total_matches_the_capped_aggregated_team_spend(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + aggregated: Final = candidate.request( + "GET", + "/team/daily/activity/aggregated", + params={**_export_range(), "team_ids": team}, + ) + assert aggregated.status_code == 200, aggregated.text + body: Final = object_value(aggregated.json()) + metadata: Final = object_value(body["metadata"]) + assert metadata["api_key_limit"] == 2, aggregated.text + assert metadata["total_api_keys"] == 3, aggregated.text + day: Final = object_value(body["results"][0]) + breakdown: Final = object_value(day["breakdown"]) + assert len(object_value(breakdown["api_keys"])) == 2, aggregated.text + team_spend: Final = float( + object_value(object_value(object_value(breakdown["entities"])[team])["metrics"])["spend"] + ) + + response: Final = _export_json(candidate, team_id=team, export_type="daily", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert len(rows) == 1, response.text + assert rows[0]["team_id"] == team, response.text + assert float(rows[0]["spend"]) == pytest.approx(team_spend), response.text + assert float(rows[0]["spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_export_users_folds_spend_per_user_and_leaves_keyless_keys_unassigned(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + user_a: Final = scenario.user(user_role="internal_user", teams=[team]) + user_b: Final = scenario.user(user_role="internal_user", teams=[team]) + key_a: Final = scenario.key(team_id=team, user_id=user_a, models=[model]) + key_b: Final = scenario.key(team_id=team, user_id=user_b, models=[model]) + key_none: Final = scenario.key(team_id=team, models=[model]) + for key in (key_a, key_b, key_none): + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_users", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + by_user: Final = {row["user_id"]: row for row in rows} + assert by_user[user_a]["spend"] == pytest.approx(spend_by_key[sha256(key_a.encode()).hexdigest()]), ( + response.text + ) + assert by_user[user_b]["spend"] == pytest.approx(spend_by_key[sha256(key_b.encode()).hexdigest()]), ( + response.text + ) + assert None in by_user, response.text + assert by_user[None]["spend"] == pytest.approx(spend_by_key[sha256(key_none.encode()).hexdigest()]), ( + response.text + ) + metadata: Final = object_value(object_value(response.json())["metadata"]) + assert float(metadata["total_spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_export_models_reports_one_row_per_model_with_matching_spend(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + upstream_a: Final = f"openai/export-{uuid.uuid4().hex}" + upstream_b: Final = f"openai/export-{uuid.uuid4().hex}" + model_a: Final = scenario.model(model=upstream_a, input_cost_per_token=0.001, output_cost_per_token=0.002) + model_b: Final = scenario.model(model=upstream_b, input_cost_per_token=0.0005, output_cost_per_token=0.001) + upstream_models: Final = (upstream_a, upstream_b) + team: Final = scenario.team(models=[model_a, model_b]) + key: Final = scenario.key(team_id=team, models=[model_a, model_b]) + for model in (model_a, model_b): + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT model, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["model"] for row in values}) == 2, + seconds=70, + ) + spend_by_model: Final = {row["model"]: float(row["spend"]) for row in daily} + + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_models", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert {row["model"] for row in rows} == set(upstream_models), response.text + for row in rows: + assert float(row["spend"]) == pytest.approx(spend_by_model[row["model"]]), response.text + + csv_response: Final = _export_json(gateway, team_id=team, export_type="daily_with_models", format="csv") + assert csv_response.status_code == 200, csv_response.text + records: Final = tuple(csv.DictReader(io.StringIO(csv_response.text))) + assert sorted(record["Model"] for record in records) == sorted(upstream_models), csv_response.text + + +def test_export_without_team_id_returns_only_the_callers_teams(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + member: Final = scenario.user(user_role="internal_user", teams=[team_a]) + member_key: Final = scenario.key(user_id=member, team_id=team_a, models=[model]) + other_key: Final = scenario.key(team_id=team_b, models=[model]) + reply: Final = gateway.chat(model, key=member_key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + reply_b: Final = gateway.chat(model, key=other_key, text=f"team export {uuid.uuid4().hex}") + assert reply_b["usage"]["total_tokens"] == 40, reply_b + eventually( + lambda: read_rows('SELECT team_id FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_b,)), + lambda values: len(values) == 1, + seconds=70, + ) + eventually( + lambda: read_rows('SELECT team_id FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_a,)), + lambda values: len(values) == 1, + seconds=70, + ) + + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert len(rows) == 1, response.text + assert rows[0]["team_id"] == team_a, response.text + + +def test_export_rejects_requests_without_a_valid_key(gateway: Gateway) -> None: + params: Final = {**_export_range(), "export_type": "daily", "format": "json"} + anonymous: Final = gateway.client.get("/team/daily/activity/export", params=params) + assert anonymous.status_code == 401, anonymous.text + garbage: Final = gateway.request("GET", "/team/daily/activity/export", params=params, key="sk-nope") + assert garbage.status_code == 401, garbage.text + + +def test_export_rejects_bad_parameters(gateway: Gateway) -> None: + weekly: Final = _export_json(gateway, export_type="weekly", format="json") + assert weekly.status_code == 422, weekly.text + xml: Final = _export_json(gateway, export_type="daily", format="xml") + assert xml.status_code == 422, xml.text + no_dates: Final = gateway.request( + "GET", "/team/daily/activity/export", params={"export_type": "daily", "format": "json"} + ) + assert no_dates.status_code == 400, no_dates.text + assert "start_date and end_date" in no_dates.text, no_dates.text + reversed_range: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={"start_date": "2026-09-25", "end_date": "2026-09-23", "export_type": "daily", "format": "json"}, + ) + assert reversed_range.status_code == 400, reversed_range.text + assert "end_date must be on or after start_date" in reversed_range.text, reversed_range.text + bad_date: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={"start_date": "2026-13-40", "end_date": "2026-12-31", "export_type": "daily", "format": "json"}, + ) + assert bad_date.status_code == 400, bad_date.text + assert "valid YYYY-MM-DD" in bad_date.text, bad_date.text + + +def test_export_of_a_team_without_spend_returns_empty(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + fresh: Final = _export_json(gateway, team_id=team, export_type="daily", format="json") + assert fresh.status_code == 200, fresh.text + body: Final = object_value(fresh.json()) + assert body["data"] == [], fresh.text + assert float(object_value(body["metadata"])["total_spend"]) == 0, fresh.text + unknown: Final = _export_json(gateway, team_id=str(uuid.uuid4()), export_type="daily", format="json") + assert unknown.status_code == 200, unknown.text + unknown_body: Final = object_value(unknown.json()) + assert unknown_body["data"] == [], unknown.text + assert float(object_value(unknown_body["metadata"])["total_spend"]) == 0, unknown.text + + +def test_export_csv_is_deterministic_and_omits_flat_cost_without_ptu(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + _team_with_three_keys(gateway, scenario, model) + params: Final = {**_export_range(), "export_type": "daily_with_keys", "format": "csv"} + first: Final = gateway.request("GET", "/team/daily/activity/export", params=params) + second: Final = gateway.request("GET", "/team/daily/activity/export", params=params) + assert first.status_code == 200 and second.status_code == 200, first.text + assert first.text == second.text, "daily_with_keys csv is not byte-identical across calls" + header: Final = first.text.splitlines()[0] + assert "Flat Cost" not in header and "Total Cost" not in header, header + + +def test_export_csv_escapes_formula_like_key_aliases(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + alias: Final = f'=HYPERLINK("http://x.{uuid.uuid4().hex}","x")' + keys: Final = ( + scenario.key(team_id=team, models=[model], key_alias=alias), + scenario.key(team_id=team, models=[model]), + ) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + eventually( + lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 2, + seconds=70, + ) + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_keys", format="csv") + assert response.status_code == 200, response.text + records: Final = {record["Key ID"]: record for record in csv.DictReader(io.StringIO(response.text))} + assert records[digests[0]]["Key Alias"] == "'" + alias, response.text + assert records[digests[1]]["Key Alias"] == "-", response.text + + +def test_aggregated_route_keeps_the_top_n_key_cap(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + response: Final = candidate.request( + "GET", + "/team/daily/activity/aggregated", + params={**_export_range(), "team_ids": team}, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + metadata: Final = object_value(body["metadata"]) + assert metadata["api_key_limit"] == 2, response.text + assert metadata["total_api_keys"] == 3, response.text + breakdown: Final = object_value(object_value(body["results"][0])["breakdown"]) + assert len(object_value(breakdown["api_keys"])) == 2, response.text + + +def test_paginated_team_daily_activity_still_lists_the_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(gateway, scenario, model) + response: Final = gateway.request( + "GET", + "/team/daily/activity", + params={ + "team_ids": team, + "start_date": _export_range()["start_date"], + "end_date": _export_range()["end_date"], + }, + ) + assert response.status_code == 200, response.text + results: Final = object_value(response.json())["results"] + assert isinstance(results, list), response.text + days: Final = tuple( + object_value(day) + for day in results + if team in object_value(object_value(object_value(day)["breakdown"])["entities"]) + ) + assert len(days) == 1, response.text + entity: Final = object_value(object_value(object_value(days[0]["breakdown"])["entities"])[team]) + assert float(object_value(entity["metrics"])["spend"]) == pytest.approx(sum(spend_by_key.values())), ( + response.text + ) + + +def test_openai_sdk_chat_still_lands_one_spend_log(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + key: Final = scenario.key(team_id=team, models=[model]) + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, max_retries=0) + reply: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"sdk {uuid.uuid4().hex}"}], stream=False + ) + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (reply.id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(rows) == 1 and rows[0]["request_id"] == reply.id, rows + + +def test_export_and_chat_burst_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + + workers: Final = eventually( + lambda: tuple(member for member in group_members(owned.process.pid) if member.pid != owned.process.pid), + lambda members: len(members) >= 2, + seconds=30, + ) + assert len(workers) >= 2, workers + + params: Final = { + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "json", + } + + def burst(tag: str) -> tuple[tuple[httpx.Response, ...], tuple[httpx.Response, ...]]: + with ThreadPoolExecutor(max_workers=30) as pool: + futures: Final = tuple( + ( + pool.submit( + candidate.request, + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"{tag}-{index}-{uuid.uuid4().hex}"}], + }, + key=keys[index % 3], + ) + if index % 2 == 0 + else pool.submit(candidate.request, "GET", "/team/daily/activity/export", params=params) + ) + for index in range(30) + ) + results: Final = tuple(future.result() for future in futures) + return results[0::2], results[1::2] + + chat_a, export_a = burst("bursta") + assert all(response.status_code == 200 for response in chat_a), [r.text for r in chat_a] + assert all(response.status_code == 200 for response in export_a), [r.text for r in export_a] + + victim: Final = workers[0] + os.kill(victim.pid, signal.SIGKILL) + + chat_b, export_b = burst("burstb") + all_chats: Final = chat_a + chat_b + all_exports: Final = export_a + export_b + assert all(response.status_code == 200 for response in all_chats), [ + (r.status_code, r.text) for r in all_chats + ] + for response in all_exports: + assert response.status_code == 200, response.text + returned: Final = {string_value(row["api_key"]) for row in object_value(response.json())["data"]} + assert returned == set(digests), response.text + chat_ids: Final = tuple(string_value(object_value(r.json())["id"]) for r in all_chats) + assert len(set(chat_ids)) == 30 + id_slots: Final = ", ".join("%s" for _ in chat_ids) + rows: Final = eventually( + lambda: read_rows( + f'SELECT request_id, COUNT(*)::int AS n FROM "LiteLLM_SpendLogs" WHERE request_id IN ({id_slots}) GROUP BY request_id', + chat_ids, + ), + lambda values: len(values) == 30, + seconds=70, + ) + assert all(row["n"] == 1 for row in rows), rows diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index 9bbc8fdde29..f85eb12c403 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -48,8 +48,9 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" "/team/daily/activity", "/team/daily/activity/aggregated", "/team/daily/activity/aggregated/search", + "/team/daily/activity/export", ), - ids=("paginated", "aggregated", "search"), + ids=("paginated", "aggregated", "search", "export"), ) @pytest.mark.parametrize( "actor,team,expected_status", @@ -59,16 +60,15 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" async def test_team_daily_activity_matrix( actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world ): + filter_param = "team_id" if endpoint.endswith("/export") else "team_ids" query = _DATES + ("&search=x" if endpoint.endswith("/search") else "") if team == "alpha": - query += f"&team_ids={world.team_alpha_id}" + query += f"&{filter_param}={world.team_alpha_id}" elif team == "beta": - query += f"&team_ids={world.team_beta_id}" + query += f"&{filter_param}={world.team_beta_id}" resp = await proxy_client.get( f"{endpoint}?{query}", headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, ) - assert ( - resp.status_code == expected_status - ), f"{actor.value} -> {team}: {resp.status_code} {resp.text}" + assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index baaf3f4ba2f..cecdb937c40 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + get_daily_activity_export_rows, global_rollup_reconciled_through, update_metrics, ) @@ -2868,3 +2869,307 @@ def test_spend_logs_window_is_none_when_no_date_parses(): from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window assert _spend_logs_window({"garbage", ""}) is None + + +_DAILY_TEAM_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyTeamSpend" ( + id TEXT PRIMARY KEY, + team_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + ptu_flat_cost DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + total_response_time_ms BIGINT DEFAULT 0, + timed_requests BIGINT DEFAULT 0 + ) +""" + + +def _seed_daily_team_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None: + with conn.cursor() as cur: + cur.execute(_DAILY_TEAM_SPEND_DDL) + cur.executemany( + """ + INSERT INTO "LiteLLM_DailyTeamSpend" + (id, team_id, date, api_key, model, model_group, custom_llm_provider, + endpoint, prompt_tokens, spend, ptu_flat_cost, api_requests, successful_requests) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + """, + rows, + ) + conn.commit() + + +def _team_spend_row( + row_id: str, + team_id: str, + api_key: str, + spend: float, + *, + date: str = "2026-06-01", + model: str = "gpt-5", + ptu_flat_cost: float = 0.0, +) -> tuple[object, ...]: + return ( + row_id, + team_id, + date, + api_key, + model, + "", + "openai", + "/v1/chat/completions", + 10, + spend, + ptu_flat_cost, + 1, + 1, + ) + + +def _export_prisma(conn: psycopg.Connection, token_rows: Sequence[SimpleNamespace] = ()) -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = _psycopg_query_raw(conn, []) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(token_rows)) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + +@pytest.mark.asyncio +async def test_export_keys_returns_every_key_beyond_the_top_n_cap( + _aggregated_postgresql: psycopg.Connection, +): + """The export route exists because the aggregated route caps the per-key arm at + USAGE_TOP_API_KEYS_LIMIT. With more keys than the cap every one of them must + land in the export, while the PTU sentinel stays out of the key view.""" + n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 7 + _seed_daily_team_spend( + _aggregated_postgresql, + [ + *[_team_spend_row(f"row-{i:03d}", "team-1", f"key-{i:03d}", float(i + 1)) for i in range(n_keys)], + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=1000.0), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_keys", + ) + + assert {row.api_key for row in rows} == {f"key-{i:03d}" for i in range(n_keys)} + assert len(rows) == n_keys + assert all(row.team_id == "team-1" for row in rows) + by_key: Final = {row.api_key: row for row in rows} + assert by_key["key-000"].spend == pytest.approx(1.0) + assert sum(row.spend for row in rows) == pytest.approx(n_keys * (n_keys + 1) / 2) + assert all(row.total_tokens == 10 and row.api_requests == 1 for row in rows) + + +@pytest.mark.asyncio +async def test_export_daily_keeps_ptu_sentinel_in_the_team_rollup( + _aggregated_postgresql: psycopg.Connection, +): + """The plain daily export groups by (date, team), so the sentinel's flat cost + must land in the team row exactly like breakdown.entities on the aggregated + route; dropping it would silently under-report team spend.""" + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=0.0), + ], + ) + with _aggregated_postgresql.cursor() as cur: + cur.execute("UPDATE \"LiteLLM_DailyTeamSpend\" SET spend = 1000.0 WHERE id = 'row-ptu'") + _aggregated_postgresql.commit() + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field={"team-1": {"team_alias": "Alpha"}}, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert len(rows) == 1 + assert rows[0].team_id == "team-1" + assert rows[0].team_alias == "Alpha" + assert rows[0].api_key is None + assert rows[0].spend == pytest.approx(1002.0) + + +@pytest.mark.asyncio +async def test_export_users_folds_keys_into_one_row_per_user( + _aggregated_postgresql: psycopg.Connection, +): + """daily_with_users runs the per-key rollup then folds in Python: two keys of + user-1 merge into one row with keys=2 and summed metrics, and the distinct + user keeps its own row.""" + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-2", "team-1", "key-2", 3.0), + _team_spend_row("row-3", "team-1", "key-3", 5.0), + ], + ) + tokens: Final = tuple( + SimpleNamespace(token=token, key_alias=None, team_id="team-1", user_id=user_id) + for token, user_id in (("key-1", "user-1"), ("key-2", "user-1"), ("key-3", "user-2")) + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql, tokens), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_users", + ) + + assert [(row.user_id, row.keys, row.spend, row.api_requests, row.total_tokens) for row in rows] == [ + ("user-1", 2, 5.0, 2, 20), + ("user-2", 1, 5.0, 1, 10), + ] + + +@pytest.mark.asyncio +async def test_export_models_rolls_up_per_team_and_model( + _aggregated_postgresql: psycopg.Connection, +): + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0, model="gpt-5"), + _team_spend_row("row-2", "team-1", "key-2", 3.0, model="gpt-5"), + _team_spend_row("row-3", "team-1", "key-1", 5.0, model="claude"), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_models", + ) + + assert [(row.model, row.spend, row.api_requests) for row in rows] == [ + ("claude", 5.0, 1), + ("gpt-5", 5.0, 2), + ] + + +@pytest.mark.asyncio +async def test_export_daily_reports_ptu_flat_cost_on_the_team_row( + _aggregated_postgresql: psycopg.Connection, ptu_cost_attribution_enabled +): + """The CSV the dashboard hands to finance must match the client-side export, + which shows flat cost columns once any PTU spend exists for the day.""" + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=240.0), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert len(rows) == 1 + assert rows[0].flat_cost == pytest.approx(240.0) + header: Final = _team_export_csv("daily", rows).splitlines()[0] + assert "Spend ($),Flat Cost ($),Total Cost ($)" in header + record: Final = _team_export_csv("daily", rows).splitlines()[1].split(",") + spend_index: Final = header.split(",").index("Spend ($)") + assert record[spend_index : spend_index + 3] == ["2.0000", "240.0000", "242.0000"] + + +@pytest.mark.asyncio +async def test_export_csv_omits_flat_cost_columns_when_no_ptu_spend_exists( + _aggregated_postgresql: psycopg.Connection, +): + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + + _seed_daily_team_spend( + _aggregated_postgresql, + [_team_spend_row("row-1", "team-1", "key-1", 2.0)], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert rows[0].flat_cost == 0.0 + header: Final = _team_export_csv("daily", rows).splitlines()[0] + assert "Flat Cost" not in header + assert "Total Cost" not in header diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 38241926f8e..6d902cb7fec 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -17101,3 +17101,130 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi assert response.status_code == 503, response.text assert response.json() == _DB_OUTAGE_503_BODY + + +def test_team_export_csv_columns_match_the_dashboard_client_layout(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias=None, + api_key="key-1", + key_alias="key-alias-1", + user_id="user-1", + user_email="u@example.com", + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + records: Final = list(csv.DictReader(io.StringIO(_team_export_csv("daily_with_keys", (row,))))) + + assert records == [ + { + "Date": "2026-06-01", + "Team": "-", + "Team ID": "team-1", + "Key Alias": "key-alias-1", + "Key ID": "key-1", + "User ID": "user-1", + "User Email": "u@example.com", + "Spend ($)": "1.5000", + "Requests": "2", + "Successful Requests": "2", + "Failed Requests": "0", + "Total Tokens": "30", + "Prompt Tokens": "20", + "Completion Tokens": "10", + "Cache Read Input Tokens": "5", + "Cache Creation Input Tokens": "4", + } + ] + + +def test_team_export_csv_omits_key_columns_for_the_plain_daily_scope(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias="Alpha", + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + text: Final = _team_export_csv("daily", (row,)) + + assert text.splitlines()[0] == ( + "Date,Team,Team ID,Spend ($),Requests,Successful Requests,Failed Requests," + "Total Tokens,Prompt Tokens,Completion Tokens,Cache Read Input Tokens,Cache Creation Input Tokens" + ) + assert list(csv.reader(io.StringIO(text)))[1] == [ + "2026-06-01", + "Alpha", + "team-1", + "1.5000", + "2", + "2", + "0", + "30", + "20", + "10", + "5", + "4", + ] + + +def test_team_export_csv_escapes_formula_aliases_and_keeps_dash_placeholder(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias='=HYPERLINK("http://evil.example","x")', + key_alias="@cmd", + user_id=None, + user_email=None, + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + record: Final = next(csv.DictReader(io.StringIO(_team_export_csv("daily_with_keys", (row,))))) + + assert record["Team"] == "'=HYPERLINK(\"http://evil.example\",\"x\")" + assert record["Key Alias"] == "'@cmd" + assert record["User ID"] == "-" + assert record["User Email"] == "-" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 16dc41c3ba8..d1ebd8b98a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -26,7 +26,7 @@ import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; -import type { EntityType } from "@/components/EntityUsageExport/types"; +import type { EntityType, ServerExport } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, customerDailyActivityCall, @@ -34,6 +34,7 @@ import { tagDailyActivityCall, teamDailyActivityAggregatedCall, teamDailyActivityCall, + teamDailyActivityExportCall, teamDailyActivityKeySearchCall, userDailyActivityCall, } from "@/components/networking"; @@ -685,7 +686,20 @@ const EntityUsage: React.FC = ({ { key: "endpoints", label: "Endpoint Activity", content: }, ]; - const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation }; + const serverExport: ServerExport | undefined = + entityType === "team" && apiKeyTruncation !== undefined && accessToken && startTime && endTime + ? (scope, format) => + teamDailyActivityExportCall({ + accessToken, + startTime, + endTime, + teamIds: entityFilterArg as string[] | null, + exportType: scope, + format, + }) + : undefined; + + const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation: serverExport ? null : apiKeyTruncation }; return (
@@ -719,6 +733,7 @@ const EntityUsage: React.FC = ({ filterOptions={getAllTags() || undefined} teams={teams || []} exportBlockedReason={getExportBlockedReason(spendFetchState)} + serverExport={serverExport} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index e0171d97423..4c5db0def62 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -253,14 +253,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { // Read through the same range stamp as the tiles, so the export is blocked from the first // render of a new range rather than from whenever the fetch effect gets around to running. + const apiKeyTruncation = getApiKeyTruncation( + userSpendData.metadata?.api_key_limit, + userSpendData.metadata?.total_api_keys, + ); const spendFetchState = { coversRange: activeAggregated !== null || paginatedResult.coversRange, cancelled: paginatedResult.cancelled, failed: paginatedResult.failed, - apiKeyTruncation: getApiKeyTruncation( - userSpendData.metadata?.api_key_limit, - userSpendData.metadata?.total_api_keys, - ), + apiKeyTruncation, }; const exportBlockedReason = getExportBlockedReason(spendFetchState); @@ -877,7 +878,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx index cb04323e4c9..6f74d241166 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx @@ -1,14 +1,3 @@ -/** - * Tests for EntityUsageExportModal component - * - * Validates core export functionality: - * - Renders modal with correct default state (CSV format, daily scope) - * - User can select export type (daily vs daily_with_models) - * - User can switch format (CSV vs JSON) - * - Export button triggers data generation with correct parameters - * - Modal closes after successful export - */ - import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -20,6 +9,7 @@ vi.mock("./utils", () => { return { handleExportCSV: vi.fn(), handleExportJSON: vi.fn(), + handleServerExport: vi.fn(async () => undefined), generateExportData: vi.fn(() => [{ Date: "2025-10-01" }]), generateMetadata: vi.fn(() => ({ meta: true })), }; @@ -114,4 +104,23 @@ describe("EntityUsageExportModal", () => { // Modal closes after export expect(baseProps.onClose).toHaveBeenCalled(); }); + + it("routes the export through the server export when one is provided, so truncated key lists still export", async () => { + /** + * When the spend fetch was capped at the top-N keys, the caller supplies a + * serverExport that hits the uncapped export route. The modal must defer to + * it instead of generating a CSV from the truncated on-screen data. + */ + const user = userEvent.setup(); + const { handleExportCSV, handleServerExport } = await import("./utils"); + const serverExport = vi.fn(async () => new Blob(["csv"])); + + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Export CSV/i })); + + expect(handleServerExport).toHaveBeenCalledWith(serverExport, "daily", "team", "csv"); + expect(handleExportCSV).not.toHaveBeenCalled(); + expect(baseProps.onClose).toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx index ec688e9fc3d..5e177a3efe6 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx @@ -10,7 +10,7 @@ import ExportFormatSelector from "./ExportFormatSelector"; import ExportSummary from "./ExportSummary"; import ExportTypeSelector from "./ExportTypeSelector"; import type { EntityUsageExportModalProps, ExportFormat, ExportScope } from "./types"; -import { handleExportCSV, handleExportJSON } from "./utils"; +import { handleExportCSV, handleExportJSON, handleServerExport } from "./utils"; const EntityUsageExportModal: React.FC = ({ isOpen, @@ -20,6 +20,7 @@ const EntityUsageExportModal: React.FC = ({ dateRange, selectedFilters, customTitle, + serverExport, }) => { const [exportFormat, setExportFormat] = useState("csv"); const [exportScope, setExportScope] = useState("daily"); @@ -35,7 +36,10 @@ const EntityUsageExportModal: React.FC = ({ const formatToUse = format || exportFormat; setIsExporting(true); try { - if (formatToUse === "csv") { + if (serverExport) { + await handleServerExport(serverExport, exportScope, entityType, formatToUse); + toast.success(`${entityLabel} usage data exported successfully as ${formatToUse.toUpperCase()}`); + } else if (formatToUse === "csv") { handleExportCSV(spendData, exportScope, entityLabel, entityType, teamAliasMap); toast.success(`${entityLabel} usage data exported successfully as CSV`); } else { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index 388a211d9bb..10825e18f3d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -16,7 +16,7 @@ import { useComboboxAnchor, } from "@/components/ui/combobox"; import EntityUsageExportModal from "./EntityUsageExportModal"; -import type { EntitySpendData, EntityType } from "./types"; +import type { EntitySpendData, EntityType, ServerExport } from "./types"; import type { Team } from "@/components/key_team_helpers/key_list"; interface UsageExportHeaderProps { @@ -35,6 +35,7 @@ interface UsageExportHeaderProps { compactLayout?: boolean; teams?: Team[]; exportBlockedReason?: string; + serverExport?: ServerExport; } const UsageExportHeader: React.FC = ({ @@ -52,6 +53,7 @@ const UsageExportHeader: React.FC = ({ compactLayout = false, teams = [], exportBlockedReason, + serverExport, }) => { const anchor = useComboboxAnchor(); const [isExportModalOpen, setIsExportModalOpen] = useState(false); @@ -142,6 +144,7 @@ const UsageExportHeader: React.FC = ({ selectedFilters={selectedFilters} customTitle={customTitle} teams={teams} + serverExport={serverExport} /> ); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts index 8491b31f5f9..3d84ef4703a 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts @@ -39,6 +39,10 @@ describe("getExportBlockedReason", () => { expect(reason).toMatch(/100 highest-spend keys of 3000/); expect(reason).toMatch(/USAGE_TOP_API_KEYS_LIMIT/); }); + + it("does not block on truncation when a server export will cover every key", () => { + expect(getExportBlockedReason(state({ apiKeyTruncation: null }))).toBeUndefined(); + }); }); describe("getApiKeyTruncation", () => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts index 6c5a5f83231..32601680bad 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts @@ -7,7 +7,7 @@ export interface UsageFetchState { coversRange: boolean; cancelled: boolean; failed: boolean; - apiKeyTruncation: ApiKeyTruncation | undefined; + apiKeyTruncation?: ApiKeyTruncation | null; } export const getApiKeyTruncation = (apiKeyLimit: unknown, totalApiKeys: unknown): ApiKeyTruncation | undefined => { @@ -25,7 +25,7 @@ export const getExportBlockedReason = ({ if (cancelled) return "Loading was stopped before the whole range arrived, so an export would under-report. Reload the page to load it all."; if (!coversRange) return "Spend data is still loading, so an export would under-report. Wait for it to finish."; - if (apiKeyTruncation !== undefined) + if (apiKeyTruncation) return `Only the ${apiKeyTruncation.limit} highest-spend keys of ${apiKeyTruncation.total} were loaded, so a per-team export would under-report. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.`; return undefined; }; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 15f193ecc3f..9cca4a2ab2b 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -17,6 +17,8 @@ export interface EntitySpendData { }; } +export type ServerExport = (exportScope: ExportScope, format: ExportFormat) => Promise; + export interface EntityUsageExportModalProps { isOpen: boolean; onClose: () => void; @@ -26,6 +28,7 @@ export interface EntityUsageExportModalProps { selectedFilters: string[]; customTitle?: string; teams?: Team[]; + serverExport?: ServerExport; } export interface ExportMetadata { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 3f9cf58ec20..4273a54e07d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -14,6 +14,7 @@ import { getEntityBreakdown, handleExportCSV, handleExportJSON, + handleServerExport, resolveEntities, } from "./utils"; @@ -3005,4 +3006,40 @@ describe("EntityUsageExport utils", () => { ]); }); }); + + describe("handleServerExport", () => { + beforeEach(() => { + document.body.innerHTML = ""; + window.URL.createObjectURL = vi.fn(() => "blob:mock-url"); + window.URL.revokeObjectURL = vi.fn(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("passes the chosen scope and format to the server export and downloads the returned blob", async () => { + const serverBlob = new Blob(["payload"], { type: "text/csv" }); + const serverExport = vi.fn(async () => serverBlob); + const createObjectURLSpy = vi.spyOn(window.URL, "createObjectURL"); + const appendChildSpy = vi.spyOn(document.body, "appendChild"); + + await handleServerExport(serverExport, "daily_with_keys", "team", "csv"); + + expect(serverExport).toHaveBeenCalledWith("daily_with_keys", "csv"); + expect(createObjectURLSpy).toHaveBeenCalledWith(serverBlob); + const attached = appendChildSpy.mock.calls[0][0] as HTMLAnchorElement; + const today = new Date().toISOString().split("T")[0]; + expect(attached.download).toBe(`team_usage_daily_with_keys_${today}.csv`); + }); + + it("lets a server failure propagate so the modal can toast it instead of downloading nothing", async () => { + const serverExport = vi.fn(async () => { + throw new Error("upstream 500"); + }); + + await expect(handleServerExport(serverExport, "daily", "team", "json")).rejects.toThrow("upstream 500"); + expect(document.body.querySelector("a")).toBeNull(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 95ce584cc89..861ea585142 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -2,21 +2,27 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import Papa from "papaparse"; import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; -import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; +import type { + EntityBreakdown, + EntitySpendData, + EntityType, + ExportFormat, + ExportMetadata, + ExportScope, + ServerExport, +} from "./types"; const resolveEntityDisplay = ( entity: string, teamAliasMap: Record, entityMetadata?: Record, -): { id: string; alias: string } => ({ - id: entity, - alias: - teamAliasMap[entity] || - entityMetadata?.team_alias || - entityMetadata?.user_email || - entityMetadata?.user_alias || - entity, -}); +): { id: string; alias: string } => { + const alias = + [teamAliasMap[entity], entityMetadata?.team_alias, entityMetadata?.user_email, entityMetadata?.user_alias].find( + Boolean, + ) ?? entity; + return { id: entity, alias }; +}; // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). // If the backend adds a field, add it here too. @@ -375,7 +381,7 @@ export const generateDailyWithModelsData = ( const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, dailyEntityMetadata[entity]); Object.entries(models).forEach(([model, metrics]: [string, any]) => { - dailyModelBreakdown.push({ + const row = { Date: day.date, [entityLabel]: alias, [`${entityLabel} ID`]: id, @@ -389,7 +395,8 @@ export const generateDailyWithModelsData = ( "Completion Tokens": metrics.completionTokens, "Cache Read Input Tokens": metrics.cacheReadInputTokens, "Cache Creation Input Tokens": metrics.cacheCreationInputTokens, - }); + }; + dailyModelBreakdown.push(row); }); }); }); @@ -449,6 +456,28 @@ export const generateMetadata = ( }; }; +export const downloadBlob = (blob: Blob, fileName: string): void => { + const url = window.URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = fileName; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + window.URL.revokeObjectURL(url); +}; + +export const handleServerExport = async ( + serverExport: ServerExport, + exportScope: ExportScope, + entityType: EntityType, + format: ExportFormat, +): Promise => { + const blob = await serverExport(exportScope, format); + const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.${format}`; + downloadBlob(blob, fileName); +}; + export const handleExportCSV = ( spendData: EntitySpendData, exportScope: ExportScope, @@ -459,15 +488,8 @@ export const handleExportCSV = ( const data = generateExportData(spendData, exportScope, entityLabel, teamAliasMap); const csv = Papa.unparse(data); const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" }); - const url = window.URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.csv`; - a.download = fileName; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - window.URL.revokeObjectURL(url); + downloadBlob(blob, fileName); }; export const handleExportJSON = ( @@ -487,13 +509,6 @@ export const handleExportJSON = ( }; const jsonString = JSON.stringify(exportObject, null, 2); const blob = new Blob([jsonString], { type: "application/json" }); - const url = window.URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.json`; - a.download = fileName; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - window.URL.revokeObjectURL(url); + downloadBlob(blob, fileName); }; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index edaf56a8a16..19af2e79bdd 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -111,6 +111,7 @@ import type { CoordinationRedisTestResponse, } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; +import type { ExportFormat, ExportScope } from "./EntityUsageExport/types"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; @@ -1467,6 +1468,36 @@ export const teamDailyActivityAggregatedCall = async ( } }; +export const teamDailyActivityExportCall = async ({ + accessToken, + startTime, + endTime, + teamIds, + exportType, + format, +}: { + accessToken: string; + startTime: Date; + endTime: Date; + teamIds: string[] | null; + exportType: ExportScope; + format: ExportFormat; +}): Promise => { + return apiClient.get(`/team/daily/activity/export`, { + accessToken, + responseType: "blob", + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + export_type: exportType, + format, + team_id: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined, + exclude_team_ids: "litellm-dashboard", + }, + }); +}; + export const teamDailyActivityKeySearchCall = async ( accessToken: string, startTime: Date, diff --git a/ui/litellm-dashboard/src/lib/http/client.ts b/ui/litellm-dashboard/src/lib/http/client.ts index c78b2b10011..974c07736e9 100644 --- a/ui/litellm-dashboard/src/lib/http/client.ts +++ b/ui/litellm-dashboard/src/lib/http/client.ts @@ -22,6 +22,8 @@ export interface RequestOptions { body?: unknown; /** Sent verbatim (FormData, Blob, pre-stringified text); disables JSON handling. */ rawBody?: BodyInit; + /** Response body handling. Defaults to JSON parsing; use this for downloads. */ + responseType?: "json" | "blob" | "text"; query?: QueryParams; headers?: Record; signal?: AbortSignal; @@ -138,7 +140,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient { const doFetch: typeof fetch = (input, init) => (fetchImpl ?? fetch)(input, init); async function request(method: HttpMethod, path: string, options: RequestOptions = {}): Promise { - const { accessToken, body, rawBody, query, headers: extraHeaders, signal, credentials } = options; + const { accessToken, body, rawBody, query, headers: extraHeaders, signal, credentials, responseType } = options; const url = appendQuery(`${getBaseUrl()}${path}`, query); @@ -177,6 +179,12 @@ export function createApiClient(config: ApiClientConfig): ApiClient { throw new ApiError(message, response.status, errorBody); } + if (responseType === "blob") { + return (await response.blob()) as T; + } + if (responseType === "text") { + return (await response.text()) as T; + } const text = await response.text(); return (text ? JSON.parse(text) : undefined) as T; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a398b2db93c..5d0bb56936a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15666,6 +15666,32 @@ export interface paths { patch?: never; trace?: never; }; + "/team/daily/activity/export": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Team Daily Activity Export + * @description Server-side Team Usage export, not subject to USAGE_TOP_API_KEYS_LIMIT. + * + * Same scoping as /team/daily/activity/aggregated, answered by one unbounded + * rollup query, returned as CSV or JSON. For daily_with_keys, + * daily_with_users and daily_with_models the PTU sentinel flat-cost rows are + * excluded, so metadata totals under those export types cover request spend + * only; the plain daily export includes them. + */ + get: operations["get_team_daily_activity_export_team_daily_activity_export_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/delete": { parameters: { query?: never; @@ -31399,7 +31425,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/export" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -42812,6 +42838,87 @@ export interface components { /** Team Id */ team_id: string; }; + /** TeamDailyActivityExportMetadata */ + TeamDailyActivityExportMetadata: { + /** End Date */ + end_date: string; + /** Export Date */ + export_date: string; + /** + * Export Type + * @enum {string} + */ + export_type: "daily" | "daily_with_keys" | "daily_with_users" | "daily_with_models"; + /** Start Date */ + start_date: string; + /** Team Ids */ + team_ids: string[] | null; + /** Total Api Requests */ + total_api_requests: number; + /** Total Failed Requests */ + total_failed_requests: number; + /** + * Total Flat Cost + * @default 0 + */ + total_flat_cost: number; + /** Total Spend */ + total_spend: number; + /** Total Successful Requests */ + total_successful_requests: number; + /** Total Tokens */ + total_tokens: number; + }; + /** TeamDailyActivityExportResponse */ + TeamDailyActivityExportResponse: { + /** Data */ + data: components["schemas"]["TeamDailyActivityExportRow"][]; + metadata: components["schemas"]["TeamDailyActivityExportMetadata"]; + }; + /** TeamDailyActivityExportRow */ + TeamDailyActivityExportRow: { + /** Api Key */ + api_key?: string | null; + /** Api Requests */ + api_requests: number; + /** Cache Creation Input Tokens */ + cache_creation_input_tokens: number; + /** Cache Read Input Tokens */ + cache_read_input_tokens: number; + /** Completion Tokens */ + completion_tokens: number; + /** Date */ + date: string; + /** Failed Requests */ + failed_requests: number; + /** + * Flat Cost + * @default 0 + */ + flat_cost: number; + /** Key Alias */ + key_alias?: string | null; + /** Keys */ + keys?: number | null; + /** Model */ + model?: string | null; + /** Prompt Tokens */ + prompt_tokens: number; + /** Spend */ + spend: number; + /** Successful Requests */ + successful_requests: number; + /** Team Alias */ + team_alias?: string | null; + /** Team Id */ + team_id: string; + /** Total Tokens */ + total_tokens: number; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id?: string | null; + }; /** * TeamListItem * @description A team item in the paginated list response, enriched with computed fields. @@ -66682,6 +66789,44 @@ export interface operations { }; }; }; + get_team_daily_activity_export_team_daily_activity_export_get: { + parameters: { + query?: { + start_date?: string | null; + end_date?: string | null; + export_type?: "daily" | "daily_with_keys" | "daily_with_users" | "daily_with_models"; + format?: "csv" | "json"; + team_id?: string | null; + exclude_team_ids?: string | null; + timezone?: number | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamDailyActivityExportResponse"]; + "text/csv": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_team_team_delete_post: { parameters: { query?: never; From fc87a06f009ea0926005b29d072f6e4d0bacf843 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:42:36 -0500 Subject: [PATCH 010/154] fix(proxy): stop leaking periodic tasks on every DB config reload (#42784) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../SlackAlerting/slack_alerting.py | 19 ++-- litellm/router.py | 6 +- .../router_strategy/base_routing_strategy.py | 16 ++- .../SlackAlerting/test_slack_alerting.py | 26 +++++ .../test_router_routing_groups.py | 105 ++++++++++++++++++ 5 files changed, 162 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8d0d044ff93..17ec3ed787d 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -118,6 +118,7 @@ class SlackAlerting(CustomBatchLogger): self.default_webhook_url = default_webhook_url self.flush_lock = asyncio.Lock() self.periodic_started = False + self._periodic_flush_task: asyncio.Task[None] | None = None self.hanging_request_check = AlertingHangingRequestCheck( slack_alerting_object=self, ) @@ -129,6 +130,12 @@ class SlackAlerting(CustomBatchLogger): self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) + def _ensure_periodic_flush_task(self) -> None: + if self.periodic_started and (self._periodic_flush_task is None or not self._periodic_flush_task.done()): + return + self._periodic_flush_task = asyncio.create_task(self.periodic_flush()) + self.periodic_started = True + def update_values( self, alerting: list | None = None, @@ -141,17 +148,14 @@ class SlackAlerting(CustomBatchLogger): ): if alerting is not None: self.alerting = alerting - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alerting_threshold is not None: self.alerting_threshold = alerting_threshold if alert_types is not None: self.alert_types = alert_types if alerting_args is not None: self.alerting_args = SlackAlertingArgs(**alerting_args) - if not self.periodic_started: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alert_type_config is not None: for key, val in alert_type_config.items(): self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val @@ -1446,9 +1450,8 @@ Model Info: return # Start periodic flush if not already started - if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + if self.alerting is not None and len(self.alerting) > 0: + self._ensure_periodic_flush_task() if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None: await self.send_webhook_alert(webhook_event=user_info) diff --git a/litellm/router.py b/litellm/router.py index 8960cd92cd8..da92eef9102 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -118,6 +118,7 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( arm_compaction, @@ -1377,6 +1378,9 @@ class Router: `_init_routing_groups`) so repeated `update_settings` calls don't accumulate dead selectors that keep receiving callback events. """ + for selector in selectors: + if isinstance(selector, BaseRoutingStrategy): + selector.retire() selector_ids: Final = {id(s) for s in selectors if s is not None} if not selector_ids: return @@ -12117,7 +12121,7 @@ class Router: ) rebuild_routing_groups = True elif var == "routing_strategy_args": - routing_args_updated = True + routing_args_updated = value != self.routing_strategy_args setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 686d57e2b77..79d457f2836 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -40,10 +40,24 @@ class BaseRoutingStrategy(ABC): self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval) ) + def cancel_sync_task(self) -> None: + if self._sync_task is not None: + self._sync_task.cancel() + + def retire(self) -> None: + self.cancel_sync_task() + if not self.redis_increment_operation_queue: + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + loop.create_task(self._push_in_memory_increments_to_redis()) + async def cleanup(self): """Cleanup method to be called when shutting down""" if self._sync_task is not None: - self._sync_task.cancel() + self.cancel_sync_task() try: await self._sync_task except asyncio.CancelledError: diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 2d5eb78950c..b9e5ff2eeb7 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -526,3 +526,29 @@ async def test_async_send_batch_collapses_only_identical_alerts() -> None: {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, {"text": CROSSED_ALERT}, ) + + +def _periodic_flush_tasks() -> list[asyncio.Task[object]]: + return [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None and t.get_coro().__qualname__ == "SlackAlerting.periodic_flush" + ] + + +@pytest.mark.asyncio +async def test_update_values_repeated_alerting_reload_keeps_single_periodic_flush_task() -> None: + slack_alerting: Final = SlackAlerting(alerting=["slack"]) + try: + for _ in range(5): + slack_alerting.update_values(alerting=["slack"]) + await asyncio.sleep(0) + flush_tasks: Final = _periodic_flush_tasks() + assert len(flush_tasks) == 1, f"expected 1 periodic_flush task, found {len(flush_tasks)}" + finally: + for t in _periodic_flush_tasks(): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 425f68dda18..534aea47885 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -18,6 +18,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -2149,3 +2150,107 @@ async def test_caller_cannot_spoof_a_priority_group_to_bypass_fallback_gates( **{metadata_bucket: {"pre_routing_selected_model": "priority-group"}}, ) assert checked == ["priority-group"] + + +def _sync_task_count() -> int: + return sum( + 1 + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ) + + +@pytest.mark.asyncio +async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tasks(monkeypatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + assert _sync_task_count() == 1 + selector_before: Final = router.lowesttpm_logger_v2 + + for _ in range(5): + router.update_settings(routing_strategy_args={"ttl": 60}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2 is selector_before, "same routing_strategy_args must not rebuild the selector" + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2.routing_args.ttl == 120 + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + + +class _RecordingRedisCache: + def __init__(self) -> None: + self.increment_lists: list[list[RedisPipelineIncrementOperation]] = [] + + async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: + self.increment_lists.append(list(increment_list)) + return [float(op["increment_value"]) for op in increment_list] + + +@pytest.mark.asyncio +async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue( + monkeypatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + redis_cache: Final = _RecordingRedisCache() + replaced: Final = router.lowesttpm_logger_v2 + replaced.dual_cache.redis_cache = redis_cache + replaced.redis_increment_operation_queue.append( + RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60) + ) + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + await asyncio.gather( + *( + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy._push_in_memory_increments_to_redis" + ) + ) + + assert router.lowesttpm_logger_v2 is not replaced + assert redis_cache.increment_lists == [ + [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)] + ] + assert replaced.redis_increment_operation_queue == [] + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass From 82146bff43d86052675e3d492661c21c17b58c01 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:43:47 -0700 Subject: [PATCH 011/154] chore(cost-map): add azure deprecation dates from the Models API for five realtime and transcribe rows (#43037) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 5 +++++ model_prices_and_context_window.json | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 90a268aa3e2..5a48b2d293b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -69776,6 +69776,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -69807,6 +69808,7 @@ "supports_tool_choice": true }, "azure/gpt-live-1": { + "deprecation_date": "2027-09-10", "input_cost_per_second": 0.000833333333333, "litellm_provider": "azure", "mode": "realtime", @@ -69824,6 +69826,7 @@ "supports_function_calling": true }, "azure/gpt-live-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 0.000283333333333, "litellm_provider": "azure", "max_input_tokens": 32000, @@ -69845,6 +69848,7 @@ "supports_audio_input": true }, "azure/gpt-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 7.5e-05, "litellm_provider": "azure", "mode": "audio_transcription", @@ -69863,6 +69867,7 @@ "supports_audio_input": true }, "azure/gpt-realtime-translate": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.000566666666667, "litellm_provider": "azure", "max_input_tokens": 32000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 90a268aa3e2..5a48b2d293b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -69776,6 +69776,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -69807,6 +69808,7 @@ "supports_tool_choice": true }, "azure/gpt-live-1": { + "deprecation_date": "2027-09-10", "input_cost_per_second": 0.000833333333333, "litellm_provider": "azure", "mode": "realtime", @@ -69824,6 +69826,7 @@ "supports_function_calling": true }, "azure/gpt-live-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 0.000283333333333, "litellm_provider": "azure", "max_input_tokens": 32000, @@ -69845,6 +69848,7 @@ "supports_audio_input": true }, "azure/gpt-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 7.5e-05, "litellm_provider": "azure", "mode": "audio_transcription", @@ -69863,6 +69867,7 @@ "supports_audio_input": true }, "azure/gpt-realtime-translate": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.000566666666667, "litellm_provider": "azure", "max_input_tokens": 32000, From 5e4b1b9df0ac766fc0640c186d994977dcf023d0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 24 Sep 2026 13:44:56 -0700 Subject: [PATCH 012/154] fix(proxy): pass team member spend rows as jsonb so a $0 flush cannot poison the pool connection (#43029) Prisma types a raw array parameter from the first batch a connection sees. After a flush in which every member cost was a whole number (a free model), the connection's cached statement expected int8[] and every later fractional batch on it failed with "improper binary format in array element", so member spend silently stopped landing while team spend kept rising. The rows now travel as one JSON document unpacked by jsonb_to_recordset with the column types declared in SQL, so Postgres types the numbers and the batch shape no longer matters. --- litellm/proxy/db/db_spend_update_writer.py | 34 +++--- .../spend/test_team_member_spend_flush.py | 106 ++++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 12 +- 3 files changed, 133 insertions(+), 19 deletions(-) create mode 100644 tests/integration/spend/test_team_member_spend_flush.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d3ac37a3e0f..17e6152bef6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -124,6 +124,12 @@ class _SpendIncrement(TypedDict): increment: ReadOnly[float] +class _MemberSpendRow(TypedDict): + user_id: ReadOnly[str] + team_id: ReadOnly[str] + cost: ReadOnly[float] + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -351,17 +357,22 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS # One statement adds every member's cost to their membership row. A missing row is created only # while the user is still on the team's roster, so a spend flush landing after a removal never -# recreates the member. +# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma +# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch +# (integers) every later fractional batch on that connection failed with "improper binary format". _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) -SELECT p.user_id, p.team_id, p.cost, p.cost -FROM unnest($1::text[], $2::text[], $3::float8[]) AS p(user_id, team_id, cost) +SELECT member.user_id, member.team_id, member.cost, member.cost +FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8) WHERE EXISTS ( SELECT 1 FROM "LiteLLM_TeamTable" t - WHERE t.team_id = p.team_id - AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id)) + WHERE t.team_id = member.team_id + AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id)) +) + OR EXISTS ( + SELECT 1 FROM "LiteLLM_TeamMembership" m + WHERE m.user_id = member.user_id AND m.team_id = member.team_id ) - OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id) ON CONFLICT (user_id, team_id) DO UPDATE SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend @@ -371,15 +382,12 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) - team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows) - for team_id in dict.fromkeys(team_ids): + for team_id in dict.fromkeys(team_id for team_id, _user_id, _cost in rows): _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) - _ = await transaction.execute_raw( - _TEAM_MEMBER_SPEND_SQL, - tuple(user_id for _team_id, user_id, _cost in rows), - team_ids, - tuple(cost for _team_id, _user_id, cost in rows), + members: Final = tuple( + _MemberSpendRow(user_id=user_id, team_id=team_id, cost=cost) for team_id, user_id, cost in rows ) + _ = await transaction.execute_raw(_TEAM_MEMBER_SPEND_SQL, json.dumps(members)) def get_llm_router(): diff --git a/tests/integration/spend/test_team_member_spend_flush.py b/tests/integration/spend/test_team_member_spend_flush.py new file mode 100644 index 00000000000..2433731f733 --- /dev/null +++ b/tests/integration/spend/test_team_member_spend_flush.py @@ -0,0 +1,106 @@ +"""Team member spend keeps landing after a flush in which every cost was a whole number. + +The proxy runs on a one-connection pool so every spend flush reuses the same database +connection. A batch of $0 requests (a free model here) is the whole-number batch, and the +fractional batches that follow it must still land on that connection. + +The $0 batch has to be flushed on its own before the paid request is sent. The spend log +row cannot prove that, since a separate monitor writes spend logs whenever they queue up, +but the daily user spend row is written by the flush cycle right after the member spend +statement, so its arrival means the whole-number batch has already been sent. +""" + +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from pydantic import JsonValue + +SINGLE_CONNECTION_CONFIG: Final = """ +model_list: [] +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 1 + proxy_batch_polling_interval: 1 + database_connection_pool_limit: 1 +router_settings: + disable_cooldowns: true +""" + + +def _member_row(team_id: str, user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', + (team_id, user_id), + ) + + +def _member_spend_is(rows: list[dict[str, JsonValue]], amount: float) -> bool: + return len(rows) == 1 and all( + float(str(rows[0][column])) == pytest.approx(amount) for column in ("spend", "total_spend") + ) + + +def _daily_user_spend_rows(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT spend FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s', (user_id,)) + + +def _logged_spend(request_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), + lambda found: len(found) == 1, + seconds=30, + ) + return float(str(rows[0]["spend"])) + + +def _chat(gateway: Gateway, key: str, model: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "member spend control"}]}, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(response.json()["id"]) + + +def test_fractional_member_spend_lands_after_a_whole_number_flush_on_the_same_connection( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = tmp_path / "single_connection_proxy.yaml" + config.write_text(SINGLE_CONNECTION_CONFIG) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + free: Final = scenario.model(input_cost_per_token=0, output_cost_per_token=0, num_retries=0) + paid: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002, num_retries=0) + team: Final = scenario.team(models=[free, paid]) + first: Final = scenario.user() + second: Final = scenario.user() + candidate.post( + "/team/member_add", + {"team_id": team, "member": [{"role": "user", "user_id": first}, {"role": "user", "user_id": second}]}, + ) + first_key: Final = scenario.key(team_id=team, user_id=first) + second_key: Final = scenario.key(team_id=team, user_id=second) + + _chat(candidate, first_key, free) + flushed: Final = eventually(lambda: _daily_user_spend_rows(first), lambda rows: len(rows) == 1, seconds=30) + assert float(str(flushed[0]["spend"])) == 0 + assert _member_spend_is(_member_row(team, first), 0) + + paid_spend: Final = _logged_spend(_chat(candidate, second_key, paid)) + assert paid_spend > 0 + eventually(lambda: _member_row(team, second), lambda rows: _member_spend_is(rows, paid_spend), seconds=30) + + repeat_spend: Final = _logged_spend(_chat(candidate, first_key, paid)) + eventually(lambda: _member_row(team, first), lambda rows: _member_spend_is(rows, repeat_spend), seconds=30) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,)), + lambda rows: float(str(rows[0]["spend"])) == pytest.approx(paid_spend + repeat_spend), + seconds=30, + ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index bf3b9aed234..7abb6e1ef92 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -994,11 +994,11 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster assert lock_statement is _TEAM_ADVISORY_LOCK_SQL assert locked_team_id == team_id assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement - statement, user_ids, team_ids, costs = spend_call.args + statement, members = spend_call.args assert statement is _TEAM_MEMBER_SPEND_SQL - assert (list(user_ids), list(team_ids), list(costs)) == ([user_id], [team_id], [response_cost]) + assert json.loads(members) == [{"user_id": user_id, "team_id": team_id, "cost": response_cost}] assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement - assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement + assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement @@ -1007,7 +1007,7 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster @pytest.mark.asyncio async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_user(): """ - The member spend statement touches rows in the order of its input arrays, so the batch + The member spend statement touches rows in the order of its input rows, so the batch is handed over sorted by (team_id, user_id), with each cost kept next to its member, and each distinct team is locked once, in `sorted(team_ids)` order, the order /team/delete locks in, so a concurrent flush and delete cannot deadlock. `eng` and `eng2` pin that: @@ -1034,13 +1034,13 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u ) *lock_calls, spend_call = mock_transaction.execute_raw.await_args_list - _statement, user_ids, team_ids, costs = spend_call.args + _statement, members = spend_call.args assert [lock_call.args for lock_call in lock_calls] == [ (_TEAM_ADVISORY_LOCK_SQL, "eng"), (_TEAM_ADVISORY_LOCK_SQL, "eng-b"), (_TEAM_ADVISORY_LOCK_SQL, "eng2"), ] - assert list(zip(team_ids, user_ids, costs)) == [ + assert [(row["team_id"], row["user_id"], row["cost"]) for row in json.loads(members)] == [ ("eng", "user_x", 0.3), ("eng", "user_y", 0.2), ("eng-b", "user_x", 0.4), From fc29fb513cf7b7cffb67018ceb968b4c20001414 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:05:31 -0700 Subject: [PATCH 013/154] fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo (#43030) * fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo Vertex only sets outputInfo.gcsOutputDirectory once a batch job has written output. Falling back to outputConfig's outputUriPrefix named the per-model directory shared by every batch of the deployment, an object that never exists, so the proxy minted a managed file for it under the first key and every other key's file calls on that id were 403s * fix(vertex_ai): treat a null gcsOutputDirectory as no output file yet --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/vertex_ai/batches/transformation.py | 27 +--- .../test_vertex_batch_output_info_wire.py | 2 +- .../vertex_ai/batches/test_transformation.py | 120 +++++++++++------- .../test_vertex_ai_batch_transformation.py | 9 +- 4 files changed, 84 insertions(+), 74 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a7dbb058465..11c99d130be 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -253,30 +253,15 @@ class VertexAIBatchTransformation: return uris[0] @classmethod - def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: + def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str | None: """ - Gets the output file id from the Vertex AI Batch response + Gets the output file id from the Vertex AI Batch response, None until Vertex reports outputInfo """ - output_info: Final = response.get("outputInfo") or OutputInfo() - output_file_id: str = output_info.get("gcsOutputDirectory", "") - if output_file_id: - output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" - if output_file_id and output_file_id != "/predictions.jsonl": - return output_file_id - - output_config: Final = response.get("outputConfig") - if output_config is None: - return output_file_id - - gcs_destination: Final = output_config.get("gcsDestination") - if gcs_destination is None: - return output_file_id - - output_uri_prefix: Final = gcs_destination.get("outputUriPrefix", "") - if output_uri_prefix.endswith("/predictions.jsonl"): - return output_uri_prefix - return output_uri_prefix.rstrip("/") + "/predictions.jsonl" + gcs_output_directory: Final = (output_info.get("gcsOutputDirectory") or "").rstrip("/") + if not gcs_output_directory: + return None + return f"{gcs_output_directory}/predictions.jsonl" @classmethod def _get_batch_job_status_from_vertex_ai_batch_response( diff --git a/tests/integration/providers/test_vertex_batch_output_info_wire.py b/tests/integration/providers/test_vertex_batch_output_info_wire.py index a7ac896076c..de41e0790d1 100644 --- a/tests/integration/providers/test_vertex_batch_output_info_wire.py +++ b/tests/integration/providers/test_vertex_batch_output_info_wire.py @@ -116,7 +116,7 @@ def test_vertex_batch_create_survives_explicit_null_output_info(gateway: Gateway "batch", "validating", _encoded(INPUT_FILE_ID, model, "file-"), - _encoded(f"{OUTPUT_PREFIX}/predictions.jsonl", model, "file-"), + None, None, "24h", ), response.text diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index ae045edbec5..5f969620a80 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -13,6 +13,8 @@ There are no real I/O seams here; ``uuid.uuid4`` is the only nondeterministic dependency and is patched where the displayName is asserted. """ +from collections.abc import Mapping +from typing import Final from unittest.mock import patch import pytest @@ -35,8 +37,7 @@ INPUT_FILE = ( ENDPOINT_ID = "7768560373388541952" ENDPOINT_INPUT_FILE = ( - f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" - "e9412502-2c91-42a6-8e61-f5c294cc0fc8" + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) @@ -248,9 +249,78 @@ def test_get_input_file_id_empty_uris(): # =========================================================================== # -# _get_output_file_id_from_vertex_ai_batch_response +# _get_output_file_id_from_vertex_ai_batch_response: None until Vertex reports outputInfo # =========================================================================== # +SHARED_OUTPUT_PREFIX: Final = "gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash" +SUCCEEDED_OUTPUT_DIRECTORY: Final = f"{SHARED_OUTPUT_PREFIX}/prediction-model-2026-09-24T19:41:00.000000Z" + + +def _vertex_job(state: str) -> dict[str, object]: + return { + "name": "projects/510528649030/locations/us-central1/batchPredictionJobs/3814889423749775360", + "state": state, + "createTime": "2026-09-24T19:37:25.775603Z", + "inputConfig": { + "instancesFormat": "jsonl", + "gcsSource": {"uris": [f"{SHARED_OUTPUT_PREFIX}/0586ba52-4f8b-4988-aa8d-3573550a4b0f"]}, + }, + "outputConfig": { + "predictionsFormat": "jsonl", + "gcsDestination": {"outputUriPrefix": SHARED_OUTPUT_PREFIX}, + }, + } + + +@pytest.mark.parametrize( + "vertex_state,output_info_field,expected_status,expected_output_file_id", + [ + ("JOB_STATE_PENDING", {}, "validating", None), + ("JOB_STATE_RUNNING", {"outputInfo": {}}, "in_progress", None), + ("JOB_STATE_CANCELLED", {"outputInfo": None}, "cancelled", None), + ( + "JOB_STATE_SUCCEEDED", + {"outputInfo": {"gcsOutputDirectory": SUCCEEDED_OUTPUT_DIRECTORY}}, + "completed", + f"{SUCCEEDED_OUTPUT_DIRECTORY}/predictions.jsonl", + ), + ], + ids=["create_or_pending", "running", "cancelled", "succeeded"], +) +def test_transform_vertex_response_output_file_id_is_none_until_output_info( + vertex_state: str, + output_info_field: Mapping[str, object], + expected_status: str, + expected_output_file_id: str | None, +) -> None: + batch: Final = T.transform_vertex_ai_batch_response_to_openai_batch_response( + {**_vertex_job(vertex_state), **output_info_field} + ) + + assert batch.status == expected_status + assert batch.output_file_id == expected_output_file_id + + +@pytest.mark.parametrize( + "response", + [ + {}, + {"outputConfig": {}}, + {"outputInfo": None}, + {"outputInfo": {"gcsOutputDirectory": ""}}, + {"outputInfo": {"gcsOutputDirectory": None}}, + ], + ids=[ + "no_fields", + "output_config_without_destination", + "null_output_info", + "empty_output_directory", + "null_output_directory", + ], +) +def test_get_output_file_id_is_none_without_output_directory(response: Mapping[str, object]) -> None: + assert T._get_output_file_id_from_vertex_ai_batch_response(response) is None + def test_get_output_file_id_from_output_info(): # outputInfo branch: rstrip trailing slash, append predictions.jsonl @@ -267,49 +337,7 @@ def test_get_output_file_id_output_info_no_trailing_slash(): ) -def test_get_output_file_id_empty_output_info_falls_through_to_output_config(): - # gcsOutputDirectory missing -> "" -> the "/predictions.jsonl" guard skips - # the outputInfo branch, falls through to outputConfig - resp = { - "outputInfo": {}, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config(): - resp = { - "outputInfo": None, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == "" - - -def test_get_output_file_id_no_output_info_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({}) == "" - - -def test_get_output_file_id_output_config_missing_gcs_destination(): - # outputConfig present but no gcsDestination -> returns the running "" value - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputConfig": {}}) == "" - - -def test_get_output_file_id_output_config_already_has_suffix(): - # outputUriPrefix already ends in /predictions.jsonl -> returned as-is (no double append) - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/predictions.jsonl"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_config_strips_trailing_slash(): - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_takes_precedence_over_output_config(): +def test_get_output_file_id_output_info_ignores_output_uri_prefix(): resp = { "outputInfo": {"gcsOutputDirectory": "gs://from-info"}, "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://from-config"}}, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 8a135dac3bb..3bc66d017a0 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -26,12 +26,12 @@ def test_output_file_id_uses_predictions_jsonl_with_output_info(): ) -def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl(): +def test_output_file_id_is_none_until_output_info(): response = { "outputInfo": {}, "outputConfig": { "gcsDestination": { - "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456" + "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro" } }, } @@ -42,10 +42,7 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl() ) ) - assert ( - output_file_id - == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl" - ) + assert output_file_id is None def test_vertex_ai_cancel_batch(): From 0f0ac4ad59cd72d8968b13e79d4e7ecf62479306 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:22:36 -0700 Subject: [PATCH 014/154] fix(playground): stop following streamed tokens, add jump to bottom button (#42968) * fix(playground): only auto-scroll the chat while pinned to the bottom Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): keep scroll pin through programmatic scrolls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): stop forcing the chat to scroll to the bottom on every update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(playground): drop scroll pinning integration tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): scroll only the chat pane, not the page, while streaming Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(playground): format ChatUI with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(playground): stop following streamed tokens, add jump to bottom button Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): jump to bottom lands on the last message, not the spacer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_ui/ChatUI.integration.test.tsx | 54 ++++++++ .../playground/components/chat_ui/ChatUI.tsx | 124 ++++++++++-------- 2 files changed, 126 insertions(+), 52 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx index e79d382ae39..a713c581e47 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx @@ -791,4 +791,58 @@ describe("ChatUI", () => { expect(screen.getByPlaceholderText("Select a Model")).toBeEnabled(); }); }); + + it("sends scroll the chat pane to the new message, tokens do not, jump button scrolls to bottom", async () => { + const scrollTopSetter = vi.spyOn(HTMLElement.prototype, "scrollTop", "set"); + let streamChunk: ((chunk: string, model?: string) => void) | undefined; + vi.mocked(makeOpenAIChatCompletionRequest).mockImplementation(async (...args) => { + streamChunk = args[1] as (chunk: string, model?: string) => void; + }); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/chat/completions"); + await selectComboboxOption("Select a Model", "Model 1"); + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + }); + await act(async () => { + fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" }); + }); + + await waitFor(() => { + expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1); + }); + + expect(scrollTopSetter).toHaveBeenCalled(); + scrollTopSetter.mockClear(); + + const scrollIntoViewMock = vi.mocked(Element.prototype.scrollIntoView); + const scrollIntoViewCallsBeforeTokens = scrollIntoViewMock.mock.calls.length; + + await act(async () => { + streamChunk?.("Hello world", "Model 1"); + }); + + expect(scrollTopSetter).not.toHaveBeenCalled(); + expect(scrollIntoViewMock.mock.calls.length).toBe(scrollIntoViewCallsBeforeTokens); + + const user = userEvent.setup(); + await user.click(screen.getByRole("button", { name: "Jump to bottom" })); + + expect(scrollTopSetter).toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index ae0fabe5ef2..a87214720ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -1,6 +1,7 @@ "use client"; import { + ArrowDown, Bot, Code2, Database, @@ -287,7 +288,7 @@ const ChatUI: React.FC = ({ // Code Interpreter state (using custom hook) const codeInterpreter = useCodeInterpreter(); - const chatEndRef = useRef(null); + const chatScrollRef = useRef(null); // Fetch MCP servers and toolsets const loadMCPServers = async () => { @@ -516,18 +517,20 @@ const ChatUI: React.FC = ({ }, [accessToken, apiKeySource, apiKey, endpointType, customProxyBaseUrl, selectedAgent]); useEffect(() => { - // Scroll to the bottom of the chat whenever chatHistory updates - if (chatEndRef.current) { - // Add a small delay to ensure content is rendered - setTimeout(() => { - chatEndRef.current?.scrollIntoView({ - behavior: "smooth", - block: "end", // Keep the scroll position at the end - }); - }, 100); - } + const el = chatScrollRef.current; + if (!el || chatHistory.at(-1)?.role !== "user") return; + const userMessages = el.querySelectorAll('[data-role="user"]'); + const last = userMessages[userMessages.length - 1]; + if (last) el.scrollTop = last.offsetTop; }, [chatHistory]); + const scrollToLastMessage = () => { + const el = chatScrollRef.current; + const messages = el?.querySelectorAll("[data-role]"); + const last = messages?.[messages.length - 1]; + if (el && last) el.scrollTop = last.offsetTop + last.offsetHeight - el.clientHeight; + }; + const handleCancelRequest = () => { if (abortControllerRef.current) { abortControllerRef.current.abort(); @@ -1801,51 +1804,68 @@ const ChatUI: React.FC = ({ )}
-
- {chatHistory.length === 0 && ( -
-
- )} - - {chatHistory.map((message, index) => ( -
- -
- ))} - - {isLoading && - mcpEvents.length > 0 && - (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && - chatHistory.length > 0 && - chatHistory[chatHistory.length - 1].role === "user" && ( -
-
-
-
-
- Assistant -
- -
+
+
+ {chatHistory.length === 0 && ( +
+
)} - {isLoading && ( -
- -
+ {chatHistory.map((message, index) => ( +
+ +
+ ))} + + {isLoading && + mcpEvents.length > 0 && + (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && + chatHistory.length > 0 && + chatHistory[chatHistory.length - 1].role === "user" && ( +
+
+
+
+
+ Assistant +
+ +
+
+ )} + + {isLoading && ( +
+ +
+ )} + {chatHistory.length > 0 &&
} +
+ {chatHistory.length > 0 && ( + )} -
From 72eb2ef651ca46f90bbbf665e99ac1c58c231d89 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:33:02 -0700 Subject: [PATCH 015/154] fix(cost-map): drop the priority input price from vertex gemini-2.5-flash-image (#43050) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 -- model_prices_and_context_window.json | 2 -- 2 files changed, 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a48b2d293b..680b6cceab6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26212,7 +26212,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -49416,7 +49415,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a48b2d293b..680b6cceab6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26212,7 +26212,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -49416,7 +49415,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, From 86ba4fc16f86731ba8736a62b1ea6dadeb2b8a48 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:35:03 -0700 Subject: [PATCH 016/154] feat(compat-matrix): resolve and install the Claude Code CLI per run (#43038) * feat(compat-matrix): resolve and install the Claude Code CLI per run * chore(compat-matrix): drop the installer header comment --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/compat-matrix-image.yml | 13 +++- tests/e2e/claude_code/cli_driver.py | 1 + tests/e2e/claude_code/cron_vm/Dockerfile | 7 -- tests/e2e/claude_code/cron_vm/README.md | 71 +++++++++++-------- .../cron_vm/install_claude_code.sh | 38 ++++++++++ tests/e2e/claude_code/cron_vm/run_daily.sh | 62 +++++++++++----- 6 files changed, 137 insertions(+), 55 deletions(-) create mode 100755 tests/e2e/claude_code/cron_vm/install_claude_code.sh diff --git a/.github/workflows/compat-matrix-image.yml b/.github/workflows/compat-matrix-image.yml index c554096cd7e..f08792c904c 100644 --- a/.github/workflows/compat-matrix-image.yml +++ b/.github/workflows/compat-matrix-image.yml @@ -4,6 +4,7 @@ on: pull_request: paths: - tests/e2e/claude_code/cron_vm/** + - tests/e2e/claude_code/pr_gate_version_resolver.py - .github/workflows/compat-matrix-image.yml workflow_dispatch: @@ -28,6 +29,14 @@ jobs: - name: Build the Render cron image run: docker build -f tests/e2e/claude_code/cron_vm/Dockerfile -t compat-matrix:${{ github.sha }} tests/e2e - - name: Run the pinned binaries as the cron user + - name: Resolve and install the Claude Code CLI as the cron user run: | - docker run --rm compat-matrix:${{ github.sha }} bash -c 'set -e; whoami; claude --version; gh --version; uv --version' + docker run --rm compat-matrix:${{ github.sha }} bash -c ' + set -euo pipefail + whoami + gh --version + uv --version + version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)" + /opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli + /tmp/claude-cli/claude --version + ' diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index a01d8ab3e7c..6996849de80 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -297,6 +297,7 @@ def run_claude( } env["ANTHROPIC_BASE_URL"] = base_url env["ANTHROPIC_AUTH_TOKEN"] = api_key + env["DISABLE_AUTOUPDATER"] = "1" # Hand the CLI a fresh empty HOME so a compromised claude package # or a model-directed Read tool call can't see the runtime user's # real dotfiles. Created here, removed in the `finally` below diff --git a/tests/e2e/claude_code/cron_vm/Dockerfile b/tests/e2e/claude_code/cron_vm/Dockerfile index 623d6b1840a..05c1387c5a2 100644 --- a/tests/e2e/claude_code/cron_vm/Dockerfile +++ b/tests/e2e/claude_code/cron_vm/Dockerfile @@ -4,8 +4,6 @@ ARG GH_VERSION=2.101.0 ARG GH_SHA256=9bca2d1c16825f109907a23307628a2f0698fbf99662b73a5cf0b020293072b8 ARG UV_VERSION=0.10.9 ARG UV_SHA256=20d79708222611fa540b5c9ed84f352bcd3937740e51aacc0f8b15b271c57594 -ARG CLAUDE_CODE_VERSION=2.1.228 -ARG CLAUDE_CODE_SHA256=d535985e6941a3eb00179ccd7f52ceb0c6623a0305a518ebc4e6514f84a94c99 SHELL ["/bin/bash", "-o", "pipefail", "-c"] @@ -23,11 +21,6 @@ RUN curl -fsSLo /tmp/uv.tar.gz "https://github.com/astral-sh/uv/releases/downloa && tar -xzf /tmp/uv.tar.gz -C /usr/local/bin --strip-components=1 uv-x86_64-unknown-linux-gnu/uv \ && rm /tmp/uv.tar.gz -RUN curl -fsSLo /tmp/claude "https://downloads.claude.ai/claude-code-releases/${CLAUDE_CODE_VERSION}/linux-x64/claude" \ - && echo "${CLAUDE_CODE_SHA256} /tmp/claude" | sha256sum -c - \ - && install -m 0755 /tmp/claude /usr/local/bin/claude \ - && rm /tmp/claude - RUN groupadd --gid 1000 populator && useradd --uid 1000 --gid 1000 --create-home populator ENV HOME=/home/populator \ diff --git a/tests/e2e/claude_code/cron_vm/README.md b/tests/e2e/claude_code/cron_vm/README.md index ed2bf4ab436..4c16b6e237a 100644 --- a/tests/e2e/claude_code/cron_vm/README.md +++ b/tests/e2e/claude_code/cron_vm/README.md @@ -16,16 +16,18 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: clone of litellm plus a cold `uv sync`. That adds a few minutes on top of the ~10 minute test run; the job's 12 hour ceiling is nowhere near. -- ⚠️ The Claude Code CLI version under test is pinned in the - `Dockerfile` (`CLAUDE_CODE_VERSION` + its checksum). Bumping it is a - PR, see the gotchas below. +- ✅ The Claude Code CLI under test is chosen on every run (the newest + npm release published at least 3 days ago) and downloaded + checksum-verified, so the matrix follows CLI releases without a PR; + see the gotchas for pinning a run. ## Layout | File | Purpose | | --- | --- | -| `Dockerfile` | The image Render builds: Debian bookworm-slim plus pinned, checksum-verified `gh`, `uv`, and the Claude Code CLI, with this `tests/e2e/` tree copied to `/opt/litellm/tests/e2e/`. Runs as the non-root user `populator` (uid/gid 1000, which is what Render's secret files are readable by). | -| `run_daily.sh` | The actual cron job. Resolves versions, clones the worktree, boots the proxy, runs pytest, builds the JSON, opens (or updates) a docs PR, sweeps stale compat-matrix PRs. | +| `Dockerfile` | The image Render builds: Debian bookworm-slim plus pinned, checksum-verified `gh` and `uv`, with this `tests/e2e/` tree copied to `/opt/litellm/tests/e2e/`. Runs as the non-root user `populator` (uid/gid 1000, which is what Render's secret files are readable by). | +| `run_daily.sh` | The actual cron job. Resolves versions, clones the worktree, installs the Claude Code CLI under test, boots the proxy, runs pytest, builds the JSON, opens (or updates) a docs PR, sweeps stale compat-matrix PRs. | +| `install_claude_code.sh` | Downloads one Claude Code release (` `) from the vendor's native release channel, verifies it against the sha256 in that release's `manifest.json`, and refuses a binary whose `--version` disagrees. Run by the cron and by the `compat-matrix-image` GitHub workflow. | | `build_matrix.py` | Tiny Python CLI that wraps `claude_code.matrix_builder.build_from_paths`. Exists only because the bash script needs *some* way to render the per-cell aggregation, and the builder is already Python. | | `check_regressions.py` | Tiny Python CLI that wraps `claude_code.matrix_builder.find_regressions`. Diffs the freshly built matrix against the currently-published one and exits `3` if any cell flipped green→red, which gates auto-merge. | | `litellm-compat-matrix.env.example` | The service's env vars, one per line, with what each is for. | @@ -35,10 +37,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: 1. **Resolves the latest LiteLLM final release tag** (newest bare `vX.Y.Z`, skipping `-rc.N`/`-dev.N` pre-releases) by paging the GitHub Releases API (`curl | jq`). -2. **Reads the Claude Code CLI version** via `claude --version`. That - is whatever the `Dockerfile` pins; the job never upgrades it on its - own. -3. **Clones the worktree** at `~/litellm-cron-worktree/` (a +2. **Clones the worktree** at `~/litellm-cron-worktree/` (a `--filter=blob:none` clone, so only the checked-out tag's blobs are fetched), `git checkout --force `, then `uv sync --frozen --no-install-project` against a uv-managed CPython 3.12 followed by @@ -52,12 +51,20 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: `claude_code/` so the tree's EKS-harness `conftest.py` (whose imports the stable venv doesn't install) is never loaded. The tag's own `tests/e2e/` is deliberately not used. +3. **Resolves and installs the Claude Code CLI under test**: + `pr_gate_version_resolver.py` (run on the venv, the image has no + Python of its own) picks the newest `@anthropic-ai/claude-code` npm + release published at least 3 days ago, the same buffer the PR gate + uses, unless `CLAUDE_CODE_VERSION` pins one, and + `install_claude_code.sh` downloads that release's `linux-x64` binary + into the run's scratch dir, verified against the release manifest. 4. **Boots the proxy** as a `setsid` background process on port `4100` bound to loopback, then polls `/health/liveliness` until it's up. -5. **Runs pytest** on `tests/e2e/claude_code/` with `LITELLM_PROXY_URL` - pointed at the proxy and `COMPAT_RESULTS_PATH` set so the conftest - hook writes the per-test results artifact. Test failures become - `fail` cells in the JSON, not script errors. +5. **Runs pytest** on `tests/e2e/claude_code/` with that CLI first on + `PATH`, `LITELLM_PROXY_URL` pointed at the proxy, and + `COMPAT_RESULTS_PATH` set so the conftest hook writes the per-test + results artifact. Test failures become `fail` cells in the JSON, not + script errors. 6. **Builds `compatibility-matrix.json`** by handing the artifact + manifest to `build_matrix.py`. 7. **Opens or updates a docs PR**: `gh repo clone` of `litellm-docs` @@ -70,7 +77,8 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: branch ... already exists" is treated as success). If the JSON is byte-identical to what `main` already publishes, the push is skipped entirely. These PRs are not gated on a second human review. -8. **Gates auto-merge on a regression check**: before enabling + + **Auto-merge is gated on a regression check**: before enabling auto-merge, `check_regressions.py` diffs the new matrix against the one currently on `main`. Auto-merge (`gh pr merge --auto --squash`) is only enabled when **no cell flipped green→red** — i.e. every @@ -82,7 +90,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: auto-merge a prior same-day run enabled is explicitly disabled — so a human reviews before it lands on the public table. The check fails *closed*: if it errors, auto-merge is withheld. -9. **Sweeps stale compat-matrix PRs**: once today's PR exists, every +8. **Sweeps stale compat-matrix PRs**: once today's PR exists, every other open `compat-matrix/*` PR that the publishing account opened from a branch on the docs repo itself is closed (and its bot-owned branch deleted), so at most one compat-matrix PR is ever open — the @@ -165,7 +173,8 @@ curl -fsS -X POST "https://api.render.com/v1/services/${CRON_ID}/deploys" \ curl -fsS "https://api.render.com/v1/services/${CRON_ID}/deploys?limit=1" \ -H "Authorization: Bearer ${RENDER_API_KEY}" -# A run that does NOT open a PR (first-time validation, CLI bumps): +# A run that does NOT open a PR (first-time validation, a CLI pinned +# with CLAUDE_CODE_VERSION): # set SKIP_PUBLISH=1 on the service, trigger a run, then remove it. # The matrix JSON is printed at the end of the run's log (nothing on # the container's disk outlives the run) and saved to @@ -217,21 +226,25 @@ docker run --rm --platform linux/amd64 \ or fine-grained Contents:RW + Pull requests:RW). It is delivered as a file, not an env var, so pytest, the proxy, and the claude CLI never inherit it; manual runs export `GITHUB_TOKEN` instead. -- **Bumping the Claude Code CLI is a PR.** Change `CLAUDE_CODE_VERSION` - in the `Dockerfile` and set `CLAUDE_CODE_SHA256` to the `linux-x64` - checksum from - `https://downloads.claude.ai/claude-code-releases//manifest.json`. - The first run on a new CLI is the riskiest one: if the new CLI - changes its wire format the matrix run can produce systematic - failures, so trigger a `SKIP_PUBLISH=1` run before the next scheduled - fire. `gh` and `uv` bump the same way, with the checksum from the - release's `gh__checksums.txt` and the tarball's `.sha256` - sidecar respectively. +- **The Claude Code CLI is chosen per run, not pinned.** Each run + tests the newest `@anthropic-ai/claude-code` npm release published + at least 3 days ago, downloaded from + `https://downloads.claude.ai/claude-code-releases//linux-x64/claude` + and verified against the sha256 in that release's `manifest.json`. + A CLI release that breaks a cell shows up as a green→red flip, which + withholds auto-merge on that day's docs PR for review. To rerun the + matrix on one specific CLI, set `CLAUDE_CODE_VERSION` on the run. + `gh` and `uv` stay pinned in the `Dockerfile`; bump them in a PR with + the checksum from the release's `gh__checksums.txt` and the + tarball's `.sha256` sidecar respectively. - **A local build on Apple silicon only proves the image assembles.** Under QEMU the Claude Code binary (a Bun executable) dies with - `CPU lacks AVX support` and `gh` panics in the Go runtime, so - `claude --version` and a full run are verified with a - `SKIP_PUBLISH=1` run on Render, not locally. + `CPU lacks AVX support` and `gh` panics in the Go runtime, so the CLI + download and `claude --version` are verified by the + `compat-matrix-image` GitHub workflow (an x86 runner that builds the + image and runs `install_claude_code.sh` in it on every PR touching + this directory) and a full run with a `SKIP_PUBLISH=1` run on Render, + not locally. - **Nothing persists between runs.** A failed run leaves no half-installed venv behind, but also no cache: don't expect a rerun to be faster than the first one. diff --git a/tests/e2e/claude_code/cron_vm/install_claude_code.sh b/tests/e2e/claude_code/cron_vm/install_claude_code.sh new file mode 100755 index 00000000000..123401cc78a --- /dev/null +++ b/tests/e2e/claude_code/cron_vm/install_claude_code.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +RELEASES_URL="https://downloads.claude.ai/claude-code-releases" + +log() { printf '==> %s\n' "$*" >&2; } +die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; } + +[[ $# -eq 2 ]] || die "usage: $(basename "$0") " +VERSION="$1" +DEST_DIR="$2" +[[ "${VERSION}" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]] || die "not a Claude Code release version: '${VERSION}'" + +mkdir -p "${DEST_DIR}" +MANIFEST="${DEST_DIR}/manifest.json" +curl -fsSL --retry 3 --retry-all-errors --output "${MANIFEST}" "${RELEASES_URL}/${VERSION}/manifest.json" \ + || die "no release manifest for claude code ${VERSION} at ${RELEASES_URL}" +CHECKSUM="$(jq -r '.platforms["linux-x64"].checksum // empty' "${MANIFEST}")" +[[ "${CHECKSUM}" =~ ^[0-9a-f]{64}$ ]] || die "manifest for claude code ${VERSION} carries no linux-x64 sha256" + +log "downloading claude code ${VERSION} (linux-x64)" +DOWNLOAD="${DEST_DIR}/claude.download" +curl -fsSL --retry 3 --retry-all-errors --output "${DOWNLOAD}" "${RELEASES_URL}/${VERSION}/linux-x64/claude" +echo "${CHECKSUM} ${DOWNLOAD}" | sha256sum -c - >/dev/null \ + || die "claude code ${VERSION} sha256 mismatch; refusing to install" +chmod 0755 "${DOWNLOAD}" +mv "${DOWNLOAD}" "${DEST_DIR}/claude" + +PROBE_HOME="$(mktemp -d -t claude-probe-home.XXXXXX)" +trap 'rm -rf "${PROBE_HOME}"' EXIT +REPORTED="$( + env -i HOME="${PROBE_HOME}" PATH="${PATH}" DISABLE_AUTOUPDATER=1 \ + "${DEST_DIR}/claude" --version | awk '{print $1}' +)" || die "claude code ${VERSION} could not run --version" +[[ "${REPORTED}" == "${VERSION}" ]] \ + || die "installed claude code reports '${REPORTED}', expected ${VERSION}" +log "installed claude code ${VERSION} at ${DEST_DIR}/claude" diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 172dd03b614..af01057eff7 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -7,15 +7,21 @@ # 1. Resolve the latest LiteLLM final release tag from the GitHub # Releases API. # 2. Update a long-lived worktree at $WORKTREE to that tag and `uv sync` it. -# 3. Boot the proxy as a background subprocess on $PROXY_PORT (default +# 3. Resolve the Claude Code CLI version under test (the newest npm +# release at least 3 days old, via pr_gate_version_resolver.py, or +# $CLAUDE_CODE_VERSION when set) and download that release's +# linux-x64 binary into the run's scratch dir, checksum-verified +# against the vendor's release manifest (install_claude_code.sh). +# 4. Boot the proxy as a background subprocess on $PROXY_PORT (default # 4100; a separate port from the human-tended :4000 proxy). -# 4. Run `pytest tests/e2e/claude_code/` against the proxy. Test -# failures become `fail` cells in the JSON, not script errors. -# 5. Hand the per-test results artifact + manifest to a small Python +# 5. Run `pytest tests/e2e/claude_code/` against the proxy with that +# CLI first on PATH. Test failures become `fail` cells in the JSON, +# not script errors. +# 6. Hand the per-test results artifact + manifest to a small Python # CLI (`build_matrix.py`) that wraps the existing # `matrix_builder.build_from_paths` to produce the published # compatibility-matrix.json. -# 6. `gh repo clone` litellm-docs, write the JSON to a deterministic +# 7. `gh repo clone` litellm-docs, write the JSON to a deterministic # branch (`compat-matrix/--`), commit, # push the branch straight to BerriAI/litellm-docs (mateo-berri has # write access), `gh pr create`, then — *only if no cell regressed @@ -23,7 +29,7 @@ # auto-merge so the PR merges itself once required checks pass. A # green→red regression leaves auto-merge off for human review; an # already-red cell (red→red) does not block. -# 7. Sweep stale compat-matrix PRs: once today's PR exists, close any +# 8. Sweep stale compat-matrix PRs: once today's PR exists, close any # other open `compat-matrix/*` PR (and delete its bot-owned branch) # so at most ONE compat-matrix PR is ever open — the newest. A # gate-withheld PR that nobody triages is superseded by the next @@ -33,7 +39,7 @@ # rather than spawning a new one. If the JSON is byte-identical to the # docs branch, we skip the push entirely. # -# Required commands on $PATH: git, uv, gh, jq, curl, claude. +# Required commands on $PATH: git, uv, gh, jq, curl. # Required state: a litellm checkout at $LITELLM_REPO (this file lives in # it); $WORKTREE is created on first run. # @@ -51,6 +57,9 @@ DOCS_BRANCH="${DOCS_BRANCH:-main}" DOCS_TARGET_PATH="${DOCS_TARGET_PATH:-src/data/compatibility-matrix.json}" SKIP_PUBLISH="${SKIP_PUBLISH:-0}" PYTEST_K="${PYTEST_K:-}" +# Empty means "resolve it": the newest @anthropic-ai/claude-code npm +# release published at least 3 days ago. Set it to pin a manual run. +CLAUDE_CODE_VERSION="${CLAUDE_CODE_VERSION:-}" # The e2e suite uses PEP 695 `type` aliases, so the venv needs Python # >= 3.12 (also what repo CI runs) even when the host's system python is # older. uv fetches a managed CPython of this version on first use -- @@ -108,7 +117,7 @@ trap cleanup EXIT INT TERM log() { printf '==> %s\n' "$*" >&2; } die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; } -for cmd in git uv gh jq curl claude; do +for cmd in git uv gh jq curl; do command -v "${cmd}" >/dev/null 2>&1 || die "missing required command: ${cmd}" done @@ -199,10 +208,6 @@ LITELLM_VERSION="$( [[ -n "${LITELLM_VERSION}" ]] || die "could not resolve latest PEP 440 final release (vX.Y.Z) in 5 pages of releases" log "resolved litellm: ${LITELLM_VERSION}" -CLAUDE_CODE_VERSION="$(claude --version 2>/dev/null | awk '{print $1}')" -[[ -n "${CLAUDE_CODE_VERSION}" ]] || die "could not read 'claude --version'" -log "local claude code: ${CLAUDE_CODE_VERSION}" - # --------------------------------------------------------------------------- # 2. Update the worktree to that tag # --------------------------------------------------------------------------- @@ -314,7 +319,27 @@ PROXY_CONFIG="${WORKTREE}/tests/e2e/claude_code/test_config.yaml" [[ -f "${PROXY_CONFIG}" ]] || die "proxy config not found at ${PROXY_CONFIG} (shim incomplete?)" # --------------------------------------------------------------------------- -# 3. Boot the proxy +# 3. Resolve and install the Claude Code CLI under test +# --------------------------------------------------------------------------- + +# The resolver is stdlib-only, but the image ships no python of its +# own, so it runs on the venv the sync above just built. Its 3-day +# publish-age buffer (PRD #26476) keeps a release that gets pulled or +# patched within days from ever driving the published matrix. +if [[ -z "${CLAUDE_CODE_VERSION}" ]]; then + CLAUDE_CODE_VERSION="$( + cd "${WORKTREE}" \ + && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../pr_gate_version_resolver.py" + )" || die "could not resolve the Claude Code version to test" + log "resolved claude code: ${CLAUDE_CODE_VERSION}" +else + log "CLAUDE_CODE_VERSION set; testing claude code ${CLAUDE_CODE_VERSION}" +fi +CLAUDE_CLI_DIR="${WORKDIR}/claude-cli" +"${POPULATOR_DIR}/install_claude_code.sh" "${CLAUDE_CODE_VERSION}" "${CLAUDE_CLI_DIR}" + +# --------------------------------------------------------------------------- +# 4. Boot the proxy # --------------------------------------------------------------------------- log "starting proxy on 127.0.0.1:${PROXY_PORT}" @@ -350,7 +375,7 @@ curl -fsS "${HEALTH_URL}" >/dev/null \ || { tail -50 "${WORKDIR}/proxy.log" >&2; die "proxy did not become healthy"; } # --------------------------------------------------------------------------- -# 4. Run pytest +# 5. Run pytest # --------------------------------------------------------------------------- RESULTS_JSON="${WORKDIR}/compat-results.json" @@ -374,6 +399,7 @@ set +e && LITELLM_PROXY_URL="http://127.0.0.1:${PROXY_PORT}" \ LITELLM_MASTER_KEY="${PROXY_API_KEY}" \ COMPAT_RESULTS_PATH="${RESULTS_JSON}" \ + PATH="${CLAUDE_CLI_DIR}:${PATH}" \ "${WORKTREE_UV}" run --no-sync pytest "${PYTEST_ARGS[@]}" ) PYTEST_EXIT=$? @@ -386,7 +412,7 @@ log "pytest exit code: ${PYTEST_EXIT} (failures become 'fail' cells, not script [[ -f "${RESULTS_JSON}" ]] || die "pytest did not produce ${RESULTS_JSON}" # --------------------------------------------------------------------------- -# 5. Build the matrix JSON +# 6. Build the matrix JSON # --------------------------------------------------------------------------- MATRIX_JSON="${WORKDIR}/compatibility-matrix.json" @@ -402,7 +428,7 @@ log "building ${MATRIX_JSON}" ) # --------------------------------------------------------------------------- -# 6. Open a docs-repo PR +# 7. Open a docs-repo PR # --------------------------------------------------------------------------- if [[ "${SKIP_PUBLISH}" == "1" ]]; then @@ -634,7 +660,9 @@ else || die "auto-merge still armed on ${BRANCH_NAME} (enabled ${AUTOMERGE_ARMED}) after --disable-auto" fi -# --- Stale-PR sweep ---------------------------------------------------------- +# --------------------------------------------------------------------------- +# 8. Sweep stale compat-matrix PRs +# --------------------------------------------------------------------------- # Keep at most ONE compat-matrix PR open: today's. Any other open # `compat-matrix/*` PR is a leftover from a day whose regression gate # withheld auto-merge and nobody triaged it; the PR we just opened or From 2be2d68ac367a84bc7035948202d59c5fa7402ea Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 14:48:35 -0700 Subject: [PATCH 017/154] test(e2e/ui): hide the LiteAdmin button in the shared admin session (#43033) #42443 pins a floating LiteAdmin button to the bottom-right corner, where it covers the logs page's next-page control. Flip the per-user Hide LiteAdmin switch during global setup so every spec reusing the admin storage state loads with the button hidden. --- tests/e2e/ui/globalSetup.ts | 4 ++++ tests/e2e/ui/helpers/navigation.ts | 14 ++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/tests/e2e/ui/globalSetup.ts b/tests/e2e/ui/globalSetup.ts index 9447c93a72e..8b3ad204767 100644 --- a/tests/e2e/ui/globalSetup.ts +++ b/tests/e2e/ui/globalSetup.ts @@ -2,6 +2,7 @@ import { chromium, expect, request } from "@playwright/test"; import { users, Role, STORAGE_PATHS } from "./fixtures/users"; import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding"; +import { hideLiteAdmin } from "./helpers/navigation"; import * as fs from "fs"; import * as path from "path"; @@ -75,6 +76,9 @@ async function globalSetup() { if (await dismiss.isVisible({ timeout: 1_500 }).catch(() => false)) { await dismiss.click(); } + if (role === Role.ProxyAdmin) { + await hideLiteAdmin(page); + } // The login flow stores a post-login return URL in the litellm_return_url // cookie. If the snapshot captures it before the app consumes it, every // test inheriting this storageState gets yanked to that stale URL the diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index e0e7b4da396..78288fab743 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -62,6 +62,20 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise } } +export async function hideLiteAdmin(page: PlaywrightPage): Promise { + await page.getByRole("button", { name: /Account menu/i }).click(); + const panel = page.getByTestId("sidebar-account-menu-panel"); + await expect(panel).toBeVisible({ timeout: 5_000 }); + const toggle = panel.getByRole("switch", { name: "Toggle hide LiteAdmin" }); + if ((await toggle.getAttribute("aria-checked")) !== "true") { + await toggle.click(); + } + await expect(toggle).toHaveAttribute("aria-checked", "true"); + await page.keyboard.press("Escape"); + await expect(panel).toBeHidden(); + await expect(page.getByRole("button", { name: "LiteAdmin", exact: true })).toBeHidden(); +} + /** * Click on a team ID in the table. Team IDs are rendered differently depending * on the component version — try button first (Tremor Button), fall back to From 1628978db707dfff9d17844d345e0587e893c307 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:49:21 -0700 Subject: [PATCH 018/154] ci: fix the litellm-tests unit job (sysmon, codecov on failure, env -i allowlist, selection errors, reruns param) (#42900) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * ci: fail the unit shard when circleci tests split errors * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/tests.yml | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 6ee1eb662e6..1afb935453f 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -74,6 +74,7 @@ commands: steps: - run: name: Install Codecov CLI (pinned v11.3.1) + when: always command: | curl -sSLf -o /tmp/codecov https://cli.codecov.io/v11.3.1/linux/codecov curl -sSLf -o /tmp/codecov.SHA256SUM https://cli.codecov.io/v11.3.1/linux/codecov.SHA256SUM @@ -90,7 +91,6 @@ commands: uv run --no-sync python -c "import litellm_enterprise; print('litellm-enterprise OK:', litellm_enterprise.__file__)" setup_test_deps: steps: - - checkout - install_uv - install_rust - restore_cache: @@ -165,9 +165,6 @@ commands: jobs: unit: parameters: - tests_path: - type: string - default: tests/unit flag: type: string default: unit @@ -180,27 +177,39 @@ jobs: pull_request_url: type: string default: "" + reruns: + type: integer + default: 0 machine: image: ubuntu-2204:2024.04.1 resource_class: large working_directory: ~/project parallelism: << parameters.shards >> environment: + COVERAGE_CORE: sysmon LITELLM_LOCAL_MODEL_COST_MAP: "True" steps: - - setup_test_deps + - checkout - skip_unless_relevant: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> + - setup_test_deps - run: - name: "Run << parameters.tests_path >> shard" + name: "Run << parameters.flag >> shard" no_output_timeout: 20m command: | mkdir -p test-results/<< parameters.flag >> - mapfile -t files < <(find << parameters.tests_path >> -name 'test_*.py' | sort | circleci tests split --split-by=timings --timings-type=filename) - if [ "${#files[@]}" -eq 0 ]; then echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.tests_path >> files; nothing to run"; exit 0; fi + selection="$(find tests/unit -name 'test_*.py' | sort)" || { echo "test selection failed for << parameters.flag >>"; exit 1; } + [ -n "${selection}" ] || { echo "test selection produced no files for << parameters.flag >>"; exit 1; } + shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; } + [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; } + mapfile -t files < <(printf '%s\n' "${shard}") + rerun_args=(-p no:rerunfailures) + if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi + test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") set +e - uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml + env -i "${test_env[@]}" \ + uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml status=$? set -e if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi @@ -224,6 +233,7 @@ jobs: resource_class: large working_directory: ~/project steps: + - checkout - setup_test_deps - run: name: Checkout litellm-docs @@ -250,16 +260,17 @@ jobs: resource_class: large working_directory: ~/project steps: - - setup_test_deps + - checkout - skip_unless_relevant: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> + - setup_test_deps - start_postgres: image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 - start_redis - run: name: Run owned integration contracts - command: bash .circleci/scripts/run_integration.sh << parameters.suite >> + command: env -i PATH="$PATH" HOME="$HOME" CIRCLE_SHA1="$CIRCLE_SHA1" CIRCLE_WORKFLOW_ID="$CIRCLE_WORKFLOW_ID" bash .circleci/scripts/run_integration.sh << parameters.suite >> no_output_timeout: 15m - run: name: Stop owned database and Redis From b72a03050140e8514ccefdab15ea9b277569ba60 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:56:52 -0700 Subject: [PATCH 019/154] test: take keys out of the legacy proxy, enterprise and mcp unit tests before moving them (#42901) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .github/workflows/test-unit-proxy-db.yml | 2 +- .../test_token_counter_gemini_contents_e2e.py | 78 ++++ .../test_prometheus_unit_tests.py | 10 +- .../routing/test_user_config_routing.py | 103 +++++ .../mcp_tests/test_aresponses_api_with_mcp.py | 377 +---------------- .../test_aresponses_api_with_mcp_providers.py | 389 ++++++++++++++++++ .../test_proxy_custom_auth.py | 5 +- .../test_proxy_pass_user_config.py | 114 ----- tests/proxy_unit_tests/test_proxy_server.py | 53 --- .../test_proxy_server_gemini_pass_through.py | 51 +++ .../test_proxy_token_counter.py | 159 +------ tests/proxy_unit_tests/test_proxy_utils.py | 4 +- 12 files changed, 652 insertions(+), 693 deletions(-) create mode 100644 tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py create mode 100644 tests/integration/routing/test_user_config_routing.py create mode 100644 tests/mcp_tests/test_aresponses_api_with_mcp_providers.py delete mode 100644 tests/proxy_unit_tests/test_proxy_pass_user_config.py create mode 100644 tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 4af7a161984..73015ac6e02 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -106,6 +106,7 @@ jobs: - test-group: proxy-server-core test-path: >- tests/proxy_unit_tests/test_proxy_server.py + tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py tests/proxy_unit_tests/test_aproxy_startup.py workers: 4 dist: loadscope @@ -115,7 +116,6 @@ jobs: tests/proxy_unit_tests/test_proxy_config_unit_test.py tests/proxy_unit_tests/test_proxy_routes.py tests/proxy_unit_tests/test_server_root_path.py - tests/proxy_unit_tests/test_proxy_pass_user_config.py tests/proxy_unit_tests/test_proxy_token_counter.py tests/proxy_unit_tests/test_request_size_limit_middleware.py tests/proxy_unit_tests/test_multipart_bypass_repro.py diff --git a/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py new file mode 100644 index 00000000000..b2b8f46ab0a --- /dev/null +++ b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py @@ -0,0 +1,78 @@ +"""Live e2e: `/utils/token_counter?call_endpoint=true` counts Gemini `contents` upstream. + +Google's countTokens API is the only tokenizer that knows Gemini's real token +boundaries, so the proxy must forward `contents` to it for both the AI Studio and +Vertex deployments and hand back the provider's `promptTokensDetails`. Claude on +Vertex is covered by `/v1/messages/count_tokens`; this is the Gemini `contents` +route the claude_code rows never reach +""" + +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call +from proxy_client import ProxyClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +GEMINI_DEPLOYMENTS = ("gemini-2.5-flash", "gemini-2.5-flash-vertex") + + +class _Part(BaseModel): + text: str + + +class _Content(BaseModel): + parts: tuple[_Part, ...] + + +class _TokenCountBody(BaseModel): + model: str + contents: tuple[_Content, ...] + + +class _CallEndpoint(BaseModel): + call_endpoint: bool = True + + +class _ModalityTokens(BaseModel): + modality: str + tokenCount: int + + +class _CountTokensUpstream(BaseModel): + totalTokens: int + promptTokensDetails: tuple[_ModalityTokens, ...] + + +class _TokenCountResponse(BaseModel): + total_tokens: int + request_model: str + model_used: str + tokenizer_type: str + original_response: _CountTokensUpstream + + +class TestGeminiContentsTokenCounting: + @pytest.mark.parametrize("model", GEMINI_DEPLOYMENTS) + def test_contents_are_counted_by_the_provider_endpoint( + self, proxy: ProxyClient, scoped_key: str, model: str + ) -> None: + text = f"Hello world, how are you doing today? {unique_marker()}" + body = _TokenCountBody(model=model, contents=(_Content(parts=(_Part(text=text),)),)) + + result = proxy.transport.send( + "/utils/token_counter", + headers=proxy.transport.bearer(scoped_key), + json=body, + params=_CallEndpoint(), + ) + + require_successful_call(result) + counted = _TokenCountResponse.model_validate_json(result.body) + assert counted.request_model == model, counted + assert counted.original_response.totalTokens == counted.total_tokens > 0, counted + assert counted.original_response.promptTokensDetails, counted + assert all(detail.tokenCount > 0 for detail in counted.original_response.promptTokensDetails), counted diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index 28fd03daf37..5b26dad269e 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -13,7 +13,6 @@ import asyncio from dotenv import load_dotenv load_dotenv() -import os from unittest.mock import MagicMock @@ -165,9 +164,9 @@ async def test_prometheus_metric_tracking(): "model_name": "gpt-5-mini", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_AI_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": "sk-azure-unit-test", + "api_version": "2025-01-01-preview", + "api_base": "https://unit-test.openai.azure.com", }, "model_info": {"id": "azure-model-id"}, }, @@ -180,9 +179,6 @@ async def test_prometheus_metric_tracking(): }, ], provider_budget_config=provider_budget_config, - redis_host=os.getenv("REDIS_HOST"), - redis_port=int(os.getenv("REDIS_PORT", 6379)), - redis_password=os.getenv("REDIS_PASSWORD"), ) try: diff --git a/tests/integration/routing/test_user_config_routing.py b/tests/integration/routing/test_user_config_routing.py new file mode 100644 index 00000000000..1c50a78201b --- /dev/null +++ b/tests/integration/routing/test_user_config_routing.py @@ -0,0 +1,103 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +USER_KEY: Final = "sk-user-supplied-" + uuid.uuid4().hex + + +def _completion(request: Request) -> Reply: + if request.target != "/v1/chat/completions": + return Reply(status=404, body=b"{}") + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-user-config", + "object": "chat.completion", + "created": 0, + "model": body["model"], + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _user_config(upstream_url: str) -> dict[str, object]: + return { + "model_list": [ + { + "model_name": "user-config-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": USER_KEY, + }, + } + ], + "num_retries": 0, + } + + +def _opt_in_config(directory: Path, upstream_url: str) -> Path: + config: Final = directory / "allow_client_side_credentials_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "admin-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": "sk-admin-configured", + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "allow_client_side_credentials": True, + }, + } + ) + ) + return config + + +def _request_body(upstream_url: str) -> dict[str, object]: + return { + "model": "user-config-deployment", + "messages": [{"role": "user", "content": "user config control"}], + "user_config": _user_config(upstream_url), + } + + +def test_user_config_routes_to_the_user_supplied_deployment_when_opted_in(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_completion) as upstream: + config: Final = _opt_in_config(tmp_path, upstream.url) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + response: Final = candidate.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "routed" + outbound: Final = tuple(upstream.received.get_nowait() for _ in range(upstream.received.qsize())) + completions: Final = tuple(request for request in outbound if request.target == "/v1/chat/completions") + assert len(completions) == 1, outbound + assert completions[0].headers["authorization"] == f"Bearer {USER_KEY}" + assert json.loads(completions[0].body)["model"] == "gpt-4.1-mini" + + +def test_user_config_is_rejected_without_the_opt_in(gateway: Gateway) -> None: + with wire_server(_completion) as upstream: + response: Final = gateway.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 401, response.text + assert "user_config is not allowed in request body" in response.text + assert upstream.received.empty() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index eb6f78b57a1..de0dc78af43 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,5 +1,3 @@ -import logging -import os import pytest from mcp.types import Tool as MCPTool from typing import List, Any, cast @@ -846,161 +844,6 @@ async def test_streaming_mcp_events_validation(): assert mock_get_tools.called, "MCP tools should have been fetched" -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - [ - pytest.param("gpt-4o-mini", id="openai"), - pytest.param("claude-haiku-4-5", id="anthropic"), - ], -) -async def test_streaming_responses_api_with_mcp_tools( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test the streaming responses API with MCP tools when using server_url="litellm_proxy" - - Under the hood the follow occurs - - - MCP: responses called litellm MCP manager.list_tools (MOCKED) - - Request 1: Made to model under test with fetched tools (REAL LLM CALL) - - MCP: Execute tool call from request 1 and returns result (MOCKED) - - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) - - Return the user the result of request 2 - """ - # Skip test if API keys are not set for the respective models - if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( - "ANTHROPIC_API_KEY" - ): - pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - print("🧪 Testing basic streaming with MCP tools...") - - # Mock MCP tools that would be returned from the manager - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "search_repo", - "description": "Search BerriAI/litellm repository for information", - "inputSchema": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"} - }, - "required": ["query"], - }, - }, by_name=False) - ] - - # Only mock the MCP-specific operations, let LLM responses be real - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - # Setup MCP mocks only - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - # Create a dynamic mock that will match the actual tool call ID from the LLM response - def mock_execute_tool_calls_side_effect( - tool_calls, user_api_key_auth, **kwargs - ): - """Mock function that returns results matching the actual tool call IDs from the LLM""" - results = [] - for tool_call in tool_calls: - # Extract call_id from the tool call - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect - - # Make the actual call - LLM responses will be real - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - tool_choice="required", - input=[ - { - "role": "user", - "type": "message", - "content": "give me a TLDR of what BerriAI/litellm is about", - } - ], - stream=True, - ) - - print(f"📋 Response type: {type(response)}") - assert hasattr( - response, "__aiter__" - ), "Response should be an async streaming response" - - # Collect streaming chunks - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") - - print(f"📊 Total chunks received: {len(chunks)}") - - # Verify MCP mocks were called (may be called multiple times in streaming) - assert ( - mock_get_tools.call_count >= 1 - ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" - print(f"MCP tools fetched: {len(mock_mcp_tools)}") - - # Verify we got a response - assert response is not None - assert len(chunks) > 0, "Should have received streaming chunks" - - print("Basic streaming responses API with MCP tools test passed!") - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) - - @pytest.mark.asyncio async def test_mcp_parameter_preparation_helpers(): """ @@ -1215,7 +1058,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): The test mocks the MCP manager response but validates the actual tools sent to the LLM to ensure no duplication occurs. """ - from unittest.mock import AsyncMock, patch, call + from unittest.mock import AsyncMock, patch from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -1432,221 +1275,3 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): "tools_per_call": [len(tools) for tools in llm_call_tools], "duplicate_tools_found": False, } - - -@pytest.mark.asyncio -@pytest.mark.parametrize("model", ["gpt-4o-mini"]) -async def test_streaming_mcp_event_order_and_response_id_consistency( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test that: - 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) - 2. All response lifecycle events share the same response ID within a cycle - """ - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "get_weather", - "description": "Get weather for a city", - "inputSchema": { - "type": "object", - "properties": { - "city": {"type": "string", "description": "City name"} - }, - "required": ["city"], - }, - }, by_name=False) - ] - - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): - results = [] - for tool_call in tool_calls: - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "Sunny, 72°F", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_side_effect - - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - input=[ - { - "role": "user", - "type": "message", - "content": "What's the weather in San Francisco?", - } - ], - stream=True, - ) - - events = [] - async for chunk in response: - events.append(chunk) - - assert len(events) > 0, "Should receive streaming events" - - created_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.created" - ), - None, - ) - in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.in_progress" - ), - None, - ) - output_item_added_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.output_item.added" - ), - None, - ) - mcp_in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) - ), - None, - ) - completed_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.completed" - ), - None, - ) - - assert created_idx is not None, "response.created event should be present" - assert ( - in_progress_idx is not None - ), "response.in_progress event should be present" - assert ( - output_item_added_idx is not None - ), "response.output_item.added event should be present" - - assert ( - created_idx < in_progress_idx - ), "response.created should come before response.in_progress" - assert ( - in_progress_idx < output_item_added_idx - ), "response.in_progress should come before response.output_item.added" - - if mcp_in_progress_idx is not None: - assert ( - output_item_added_idx < mcp_in_progress_idx - ), "response.output_item.added should come before response.mcp_list_tools.in_progress" - - response_ids = [] - for i, event in enumerate(events): - event_type = getattr(event, "type", None) - if hasattr(event, "response"): - response_obj = getattr(event, "response", None) - if response_obj and hasattr(response_obj, "id"): - event_type_value = ( - event_type.value - if hasattr(event_type, "value") - else str(event_type) - ) - if any( - x in event_type_value - for x in [ - "response.created", - "response.in_progress", - "response.completed", - ] - ): - response_ids.append((i, event_type_value, response_obj.id)) - - assert ( - len(response_ids) >= 2 - ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" - - cycles = [] - current_cycle = [] - current_id = None - - for idx, event_type, resp_id in response_ids: - if current_id is None or resp_id == current_id: - current_cycle.append((idx, event_type, resp_id)) - current_id = resp_id - else: - if current_cycle: - cycles.append(current_cycle) - current_cycle = [(idx, event_type, resp_id)] - current_id = resp_id - if current_cycle: - cycles.append(current_cycle) - - for cycle_num, cycle in enumerate(cycles): - cycle_ids = set(resp_id for _, _, resp_id in cycle) - assert ( - len(cycle_ids) == 1 - ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" - - assert ( - completed_idx is not None - ), "response.completed event should be present" - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py new file mode 100644 index 00000000000..72a0415cf88 --- /dev/null +++ b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py @@ -0,0 +1,389 @@ +import logging +import os +import pytest +from mcp.types import Tool as MCPTool +from typing import Any, cast + +import litellm +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + +class MockUserAPIKeyAuth: + """Mock UserAPIKeyAuth for testing""" + + def __init__(self): + self.api_key = "test_key" + self.user_id = "test_user" + self.team_id = "test_team" + self.user_email = "test@example.com" + self.max_budget = 100.0 + self.spend = 0.0 + self.models = [] + self.aliases = {} + self.config = {} + self.permissions = {} + self.metadata = {} + self.object_permission_id = "test_permission_id" + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + pytest.param("gpt-4o-mini", id="openai"), + pytest.param("claude-haiku-4-5", id="anthropic"), + ], +) +async def test_streaming_responses_api_with_mcp_tools( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test the streaming responses API with MCP tools when using server_url="litellm_proxy" + + Under the hood the follow occurs + + - MCP: responses called litellm MCP manager.list_tools (MOCKED) + - Request 1: Made to model under test with fetched tools (REAL LLM CALL) + - MCP: Execute tool call from request 1 and returns result (MOCKED) + - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) + + Return the user the result of request 2 + """ + if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( + "ANTHROPIC_API_KEY" + ): + pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + print("🧪 Testing basic streaming with MCP tools...") + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "search_repo", + "description": "Search BerriAI/litellm repository for information", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_tool_calls_side_effect( + tool_calls, user_api_key_auth, **kwargs + ): + """Mock function that returns results matching the actual tool call IDs from the LLM""" + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + tool_choice="required", + input=[ + { + "role": "user", + "type": "message", + "content": "give me a TLDR of what BerriAI/litellm is about", + } + ], + stream=True, + ) + + print(f"📋 Response type: {type(response)}") + assert hasattr( + response, "__aiter__" + ), "Response should be an async streaming response" + + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") + + print(f"📊 Total chunks received: {len(chunks)}") + + assert ( + mock_get_tools.call_count >= 1 + ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" + print(f"MCP tools fetched: {len(mock_mcp_tools)}") + + assert response is not None + assert len(chunks) > 0, "Should have received streaming chunks" + + print("Basic streaming responses API with MCP tools test passed!") + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +async def test_streaming_mcp_event_order_and_response_id_consistency( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test that: + 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) + 2. All response lifecycle events share the same response ID within a cycle + """ + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "get_weather", + "description": "Get weather for a city", + "inputSchema": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "Sunny, 72°F", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + input=[ + { + "role": "user", + "type": "message", + "content": "What's the weather in San Francisco?", + } + ], + stream=True, + ) + + events = [] + async for chunk in response: + events.append(chunk) + + assert len(events) > 0, "Should receive streaming events" + + created_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.created" + ), + None, + ) + in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.in_progress" + ), + None, + ) + output_item_added_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.output_item.added" + ), + None, + ) + mcp_in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) + ), + None, + ) + completed_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.completed" + ), + None, + ) + + assert created_idx is not None, "response.created event should be present" + assert ( + in_progress_idx is not None + ), "response.in_progress event should be present" + assert ( + output_item_added_idx is not None + ), "response.output_item.added event should be present" + + assert ( + created_idx < in_progress_idx + ), "response.created should come before response.in_progress" + assert ( + in_progress_idx < output_item_added_idx + ), "response.in_progress should come before response.output_item.added" + + if mcp_in_progress_idx is not None: + assert ( + output_item_added_idx < mcp_in_progress_idx + ), "response.output_item.added should come before response.mcp_list_tools.in_progress" + + response_ids = [] + for i, event in enumerate(events): + event_type = getattr(event, "type", None) + if hasattr(event, "response"): + response_obj = getattr(event, "response", None) + if response_obj and hasattr(response_obj, "id"): + event_type_value = ( + event_type.value + if hasattr(event_type, "value") + else str(event_type) + ) + if any( + x in event_type_value + for x in [ + "response.created", + "response.in_progress", + "response.completed", + ] + ): + response_ids.append((i, event_type_value, response_obj.id)) + + assert ( + len(response_ids) >= 2 + ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" + + cycles = [] + current_cycle = [] + current_id = None + + for idx, event_type, resp_id in response_ids: + if current_id is None or resp_id == current_id: + current_cycle.append((idx, event_type, resp_id)) + current_id = resp_id + else: + if current_cycle: + cycles.append(current_cycle) + current_cycle = [(idx, event_type, resp_id)] + current_id = resp_id + if current_cycle: + cycles.append(current_cycle) + + for cycle_num, cycle in enumerate(cycles): + cycle_ids = set(resp_id for _, _, resp_id in cycle) + assert ( + len(cycle_ids) == 1 + ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" + + assert ( + completed_idx is not None + ), "response.completed event should be present" + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) diff --git a/tests/proxy_unit_tests/test_proxy_custom_auth.py b/tests/proxy_unit_tests/test_proxy_custom_auth.py index b575e4c85c6..dbbad0dab1e 100644 --- a/tests/proxy_unit_tests/test_proxy_custom_auth.py +++ b/tests/proxy_unit_tests/test_proxy_custom_auth.py @@ -53,8 +53,7 @@ def test_custom_auth(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") - print(f"token: {token}") + token = "sk-unit-test-master" headers = {"Authorization": f"Bearer {token}"} with pytest.raises(Exception, match="Authentication Error, Failed custom auth") as exc_info: client.post("/chat/completions", json=test_data, headers=headers) @@ -71,7 +70,7 @@ def test_custom_auth_bearer(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") + token = "sk-unit-test-master" headers = {"Authorization": f"WITHOUT BEAR Er {token}"} with pytest.raises(Exception, match="CustomAuth - Malformed API Key passed in") as exc_info: diff --git a/tests/proxy_unit_tests/test_proxy_pass_user_config.py b/tests/proxy_unit_tests/test_proxy_pass_user_config.py deleted file mode 100644 index 91911c142ea..00000000000 --- a/tests/proxy_unit_tests/test_proxy_pass_user_config.py +++ /dev/null @@ -1,114 +0,0 @@ -import sys, os -import traceback -from dotenv import load_dotenv - -load_dotenv() -import io - -# this file is to test litellm/proxy - -import pytest, logging, asyncio -import litellm -from litellm import embedding, completion, completion_cost, Timeout -from litellm import RateLimitError - -# Configure logging -logging.basicConfig( - level=logging.DEBUG, # Set the desired logging level - format="%(asctime)s - %(levelname)s - %(message)s", -) - -# test /chat/completion request to the proxy -from fastapi.testclient import TestClient -from fastapi import FastAPI -from litellm.proxy.proxy_server import ( - router, - save_worker_config, - initialize, -) # Replace with the actual module where your FastAPI router is defined - -# Your bearer token -token = "sk-1234" - -headers = {"Authorization": f"Bearer {token}"} - - -@pytest.fixture(scope="function") -def client_no_auth(): - # Assuming litellm.proxy.proxy_server is an object - from litellm.proxy.proxy_server import cleanup_router_config_variables - - cleanup_router_config_variables() - filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" - # initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables - asyncio.run(initialize(config=config_fp, debug=True)) - app = FastAPI() - app.include_router(router) # Include your router in the test app - - return TestClient(app) - - -@pytest.mark.skipif( - os.environ.get("AZURE_AI_API_KEY") is None - or os.environ.get("OPENAI_API_KEY") is None, - reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test", -) -def test_chat_completion(client_no_auth): - global headers - - from litellm.types.router import RouterConfig, ModelConfig - from litellm.types.completion import CompletionRequest - - user_config = RouterConfig( - model_list=[ - ModelConfig( - model_name="user-azure-instance", - litellm_params=CompletionRequest( - model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_AI_API_KEY"), - api_version=os.getenv("AZURE_API_VERSION"), - api_base=os.getenv("AZURE_AI_API_BASE"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ModelConfig( - model_name="user-openai-instance", - litellm_params=CompletionRequest( - model="gpt-3.5-turbo", - api_key=os.getenv("OPENAI_API_KEY"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ], - num_retries=2, - allowed_fails=3, - fallbacks=[{"user-azure-instance": ["user-openai-instance"]}], - ).dict() - - try: - # Your test data - test_data = { - "model": "user-azure-instance", - "messages": [ - {"role": "user", "content": "hi"}, - ], - "max_tokens": 10, - "user_config": user_config, - } - - print("testing proxy server with chat completions") - response = client_no_auth.post("/v1/chat/completions", json=test_data) - print(f"response - {response.text}") - assert response.status_code == 200 - result = response.json() - print(f"Received response: {result}") - except Exception as e: - pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") - - -# Run the test diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index ed0380058a5..5be27b3ad72 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2065,59 +2065,6 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( assert new_data["failure_callback"] == expected_failure_callbacks -@pytest.mark.skipif( - not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), - reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", -) -@pytest.mark.asyncio -async def test_gemini_pass_through_endpoint(): - from starlette.datastructures import URL - - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - Request, - Response, - gemini_proxy_route, - ) - - body = b""" - { - "contents": [{ - "parts":[{ - "text": "The quick brown fox jumps over the lazy dog." - }] - }] - } - """ - - # Construct the scope dictionary - scope = { - "type": "http", - "method": "POST", - "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", - "query_string": b"key=sk-1234", - "headers": [ - (b"content-type", b"application/json"), - ], - } - - # Create a new Request object - async def async_receive(): - return {"type": "http.request", "body": body, "more_body": False} - - request = Request( - scope=scope, - receive=async_receive, - ) - - resp = await gemini_proxy_route( - endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", - request=request, - fastapi_response=Response(), - ) - - print(resp.body) - - @pytest.mark.parametrize("hidden", [True, False]) @pytest.mark.asyncio async def test_model_info_alias_without_prisma(hidden): diff --git a/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py new file mode 100644 index 00000000000..2453ec3bfe3 --- /dev/null +++ b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py @@ -0,0 +1,51 @@ +import os + +import pytest + + +@pytest.mark.skipif( + not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), + reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", +) +@pytest.mark.asyncio +async def test_gemini_pass_through_endpoint(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + Request, + Response, + gemini_proxy_route, + ) + + body = b""" + { + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] + } + """ + + scope = { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", + "query_string": b"key=sk-1234", + "headers": [ + (b"content-type", b"application/json"), + ], + } + + async def async_receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + scope=scope, + receive=async_receive, + ) + + await gemini_proxy_route( + endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", + request=request, + fastapi_response=Response(), + ) + diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39ec4bb1887..8590e959961 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -2,10 +2,7 @@ # 1. Generate a Key, and use it to make a call -import json import logging -import os -import tempfile from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -35,79 +32,6 @@ from litellm.types.utils import TokenCountResponse verbose_proxy_logger.setLevel(level=logging.DEBUG) -def get_vertex_ai_creds_json() -> dict: - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - return service_account_key_data - - -def load_vertex_ai_credentials(): - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) - - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -223,10 +147,12 @@ async def test_anthropic_messages_count_tokens_endpoint(): - Should return response in Anthropic format: {"input_tokens": } - Should work as wrapper around internal token_counter function """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -295,10 +221,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): - Should still work and return Anthropic format - Should call internal token_counter with from_anthropic_endpoint=True """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -435,10 +363,12 @@ async def test_anthropic_endpoint_error_handling(): """ Test error handling in the /v1/messages/count_tokens endpoint """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request, HTTPException from unittest.mock import MagicMock + from fastapi import HTTPException, Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_user_api_key_dict = MagicMock() @@ -474,8 +404,10 @@ async def test_anthropic_endpoint_error_handling(): @pytest.mark.asyncio async def test_factory_anthropic_endpoint_calls_anthropic_counter(): """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -531,8 +463,10 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): @pytest.mark.asyncio async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -590,8 +524,10 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): @pytest.mark.asyncio async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -678,57 +614,6 @@ async def test_factory_registration(): assert not counter.should_use_token_counting_api(custom_llm_provider=None) -@pytest.mark.skip( - reason="Requires Google/Vertex AI credentials (GEMINI_API_KEY or VERTEX_AI_PRIVATE_KEY)." -) -@pytest.mark.asyncio -@pytest.mark.parametrize("model_name", ["gemini-2.5-pro", "vertex-ai-gemini-2.5-pro"]) -async def test_vertex_ai_gemini_token_counting_with_contents(model_name): - """ - Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True - """ - load_vertex_ai_credentials() - llm_router = Router( - model_list=[ - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "gemini/gemini-2.5-pro", - }, - }, - { - "model_name": "vertex-ai-gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - }, - }, - ] - ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with contents format and call_endpoint=True - response = await token_counter( - request=TokenCountRequest( - model=model_name, - contents=[ - {"parts": [{"text": "Hello world, how are you doing today? i am ij"}]} - ], - ), - call_endpoint=True, - ) - - print("Vertex AI Gemini token counting response:", response) - - # validate we have original response - assert response.original_response is not None - assert response.original_response.get("totalTokens") is not None - assert response.original_response.get("promptTokensDetails") is not None - - prompt_tokens_details = response.original_response.get("promptTokensDetails") - assert prompt_tokens_details is not None - - @pytest.mark.asyncio async def test_bedrock_count_tokens_endpoint(): """ @@ -779,7 +664,7 @@ async def test_vertex_ai_anthropic_token_counting(): This tests the token counting implementation for Vertex AI partner models without making actual API calls. Mocks at the handler level to test the full flow. """ - from unittest.mock import AsyncMock, patch, MagicMock + from unittest.mock import patch # Mock the Vertex AI partner models token counter response mock_token_response = { diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 1134f41a940..ab787bbbe27 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -438,7 +438,7 @@ def test_is_request_body_safe_global_enabled( "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), + "api_key": "sk-openai-unit-test", }, } ] @@ -475,7 +475,7 @@ def test_is_request_body_safe_model_enabled( "model_name": "fireworks_ai/*", "litellm_params": { "model": "fireworks_ai/*", - "api_key": os.getenv("FIREWORKS_API_KEY"), + "api_key": "sk-fireworks-unit-test", "configurable_clientside_auth_params": ( ["api_base"] if allow_client_side_credentials else [] ), From e3f087315de3eac8ec5c78b31fca218b1f846892 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:06:30 -0500 Subject: [PATCH 020/154] feat(terraform): add display_name to litellm_model resource and model data sources (#42987) * feat(terraform): add display_name to litellm_model resource and model data sources Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): persist display_name on update and read /model/info data envelope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(terraform): drop PATCH /model/{model_id}/update from endpoint audit allowlist Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): surface external display_name removal as drift on refresh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(terraform): rerun after uv download timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 2 + terraform/provider/docs/data-sources/model.md | 1 + .../provider/docs/data-sources/models.md | 1 + terraform/provider/docs/resources/model.md | 2 + .../provider/litellm/data_source_model.go | 24 +- .../litellm/data_source_model_test.go | 12 +- terraform/provider/litellm/resource_model.go | 5 + .../provider/litellm/resource_model_crud.go | 35 ++- .../provider/litellm/resource_model_test.go | 244 ++++++++++++++++++ terraform/provider/litellm/types.go | 23 +- terraform/provider/litellm/utils.go | 7 + .../endpointaudit/coverage_allowlist.txt | 1 - 12 files changed, 333 insertions(+), 24 deletions(-) create mode 100644 terraform/provider/litellm/resource_model_test.go diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 079bb7d8667..4123b7e6b51 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources - **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement @@ -48,6 +49,7 @@ longer signal it. ### Fixed +- **model**: `litellm_model` refresh now reads the `{"data": [...]}` envelope `/model/info` returns, so `model_info` fields changed outside Terraform show up as drift instead of silently keeping the previous state - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update - **credential**: create now reports a `credential_name` collision as a clear error naming the `terraform import` command that adopts the existing credential, instead of surfacing the proxy's raw 500 with a Prisma `Unique constraint failed` message. New `adopt_existing` argument (default `false`) opts into taking the existing credential over during create, which makes `apply` idempotent again once state loses track of a credential that still exists on the proxy. Requires a proxy that answers 409 on the collision; older proxies are still detected by their 500 message - **credential**: credential names and `model_id` are now percent-encoded in request URLs, so a name containing `/`, `?`, `#` or spaces reaches the proxy intact instead of being cut at the first reserved character and read, updated or deleted as a different credential diff --git a/terraform/provider/docs/data-sources/model.md b/terraform/provider/docs/data-sources/model.md index 6976ff1523a..587652bcbd8 100644 --- a/terraform/provider/docs/data-sources/model.md +++ b/terraform/provider/docs/data-sources/model.md @@ -43,6 +43,7 @@ In addition to all arguments above, the following attributes are exported: * `tier` - Model tier (`free` or `paid`). * `mode` - Model mode, e.g. `chat` or `embedding`. * `team_id` - Team the deployment is scoped to, if any. +* `display_name` - Human-readable name returned by `/v1/models`, if configured. * `db_model` - Whether the deployment is stored in the database (as opposed to config). ## Security Note diff --git a/terraform/provider/docs/data-sources/models.md b/terraform/provider/docs/data-sources/models.md index 7862dc30ab7..1cf0fd36ab3 100644 --- a/terraform/provider/docs/data-sources/models.md +++ b/terraform/provider/docs/data-sources/models.md @@ -41,4 +41,5 @@ In addition to all arguments above, the following attributes are exported: * `tier` - Model tier (`free` or `paid`). * `mode` - Model mode, e.g. `chat` or `embedding`. * `team_id` - Team the deployment is scoped to, if any. + * `display_name` - Human-readable name returned by `/v1/models`, if configured. * `db_model` - Whether the deployment is stored in the database. diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 0409b48b391..5bb68bfe918 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -126,6 +126,8 @@ The following arguments are supported: * `team_id` - (Optional) string. Associate the model with a specific team. +* `display_name` - (Optional) string. Human-readable name stored in `model_info.display_name` and returned as `display_name` by `/v1/models`, so clients such as Claude Code and Claude Desktop show it in their model picker instead of `model_name`. When unset, clients fall back to `model_name`. + * `mode` - (Optional) string. The intended use of the model. Valid values are: * `completion` * `embedding` diff --git a/terraform/provider/litellm/data_source_model.go b/terraform/provider/litellm/data_source_model.go index 78af04ac160..6b9993d267e 100644 --- a/terraform/provider/litellm/data_source_model.go +++ b/terraform/provider/litellm/data_source_model.go @@ -23,14 +23,15 @@ type modelInfoParams struct { } type modelInfoMeta struct { - ID string `json:"id"` - DBModel bool `json:"db_model"` - BaseModel string `json:"base_model"` - Tier string `json:"tier"` - Mode string `json:"mode"` - TeamID string `json:"team_id"` - CreatedAt string `json:"created_at"` - UpdatedAt string `json:"updated_at"` + ID string `json:"id"` + DBModel bool `json:"db_model"` + BaseModel string `json:"base_model"` + Tier string `json:"tier"` + Mode string `json:"mode"` + TeamID string `json:"team_id"` + DisplayName string `json:"display_name"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` } type modelInfoEntry struct { @@ -112,6 +113,10 @@ func dataSourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Computed: true, }, + "display_name": { + Type: schema.TypeString, + Computed: true, + }, "db_model": { Type: schema.TypeBool, Computed: true, @@ -161,6 +166,7 @@ func dataSourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("tier", entry.ModelInfo.Tier) d.Set("mode", entry.ModelInfo.Mode) d.Set("team_id", entry.ModelInfo.TeamID) + d.Set("display_name", entry.ModelInfo.DisplayName) d.Set("db_model", entry.ModelInfo.DBModel) log.Printf("[INFO] Successfully read model with ID: %s", modelID) @@ -197,6 +203,7 @@ func dataSourceLiteLLMModels() *schema.Resource { "tier": {Type: schema.TypeString, Computed: true}, "mode": {Type: schema.TypeString, Computed: true}, "team_id": {Type: schema.TypeString, Computed: true}, + "display_name": {Type: schema.TypeString, Computed: true}, "db_model": {Type: schema.TypeBool, Computed: true}, }, }, @@ -247,6 +254,7 @@ func dataSourceLiteLLMModelsRead(d *schema.ResourceData, m interface{}) error { "tier": entry.ModelInfo.Tier, "mode": entry.ModelInfo.Mode, "team_id": entry.ModelInfo.TeamID, + "display_name": entry.ModelInfo.DisplayName, "db_model": entry.ModelInfo.DBModel, }) } diff --git a/terraform/provider/litellm/data_source_model_test.go b/terraform/provider/litellm/data_source_model_test.go index 97d7f07dcd8..b46a355853a 100644 --- a/terraform/provider/litellm/data_source_model_test.go +++ b/terraform/provider/litellm/data_source_model_test.go @@ -35,7 +35,8 @@ func TestDataSourceModelReadSingleObject(t *testing.T) { "base_model": "gpt-4o", "tier": "paid", "mode": "chat", - "team_id": "team-1" + "team_id": "team-1", + "display_name": "GPT-4o" } } }`)) @@ -66,6 +67,7 @@ func TestDataSourceModelReadSingleObject(t *testing.T) { "tier": "paid", "mode": "chat", "team_id": "team-1", + "display_name": "GPT-4o", "db_model": true, } for attr, want := range checks { @@ -115,7 +117,7 @@ func TestDataSourceModelsRead(t *testing.T) { w.Header().Set("Content-Type", "application/json") w.Write([]byte(`{ "data": [ - {"model_name": "a", "litellm_params": {"model": "openai/a", "custom_llm_provider": "openai"}, "model_info": {"id": "id-1", "db_model": true}}, + {"model_name": "a", "litellm_params": {"model": "openai/a", "custom_llm_provider": "openai"}, "model_info": {"id": "id-1", "db_model": true, "display_name": "Model A"}}, {"model_name": "b", "litellm_params": {"model": "anthropic/b", "custom_llm_provider": "anthropic"}, "model_info": {"id": "id-2"}} ] }`)) @@ -143,7 +145,11 @@ func TestDataSourceModelsRead(t *testing.T) { t.Fatalf("expected 2 models, got %d", len(models)) } first := models[0].(map[string]interface{}) - if first["model_name"] != "a" || first["custom_llm_provider"] != "openai" || first["db_model"] != true { + if first["model_name"] != "a" || first["custom_llm_provider"] != "openai" || first["db_model"] != true || first["display_name"] != "Model A" { t.Errorf("unexpected first model: %v", first) } + second := models[1].(map[string]interface{}) + if second["display_name"] != "" { + t.Errorf("expected empty display_name for model without one, got %v", second["display_name"]) + } } diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index b0a7304718b..85cb1d038bf 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -93,6 +93,11 @@ func resourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Optional: true, }, + "display_name": { + Type: schema.TypeString, + Optional: true, + Description: "Human-readable name returned as display_name by /v1/models, shown in client model pickers instead of model_name", + }, "mode": { Type: schema.TypeString, Optional: true, diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index fc5d5b09dd5..c7db2a673e5 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "log" + "net/url" "strconv" "strings" "time" @@ -53,6 +54,7 @@ func retryModelRead(d *schema.ResourceData, m interface{}, maxRetries int) error const ( endpointModelNew = "/model/new" endpointModelUpdate = "/model/update" + endpointModelPatch = "/model/%s/update" endpointModelInfo = "/model/info" endpointModelDelete = "/model/delete" ) @@ -246,12 +248,13 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e ModelName: d.Get("model_name").(string), LiteLLMParams: litellmParams, ModelInfo: ModelInfo{ - ID: modelID, - DBModel: true, - BaseModel: pricingBaseModel, - Tier: d.Get("tier").(string), - Mode: d.Get("mode").(string), - TeamID: d.Get("team_id").(string), + ID: modelID, + DBModel: true, + BaseModel: pricingBaseModel, + Tier: d.Get("tier").(string), + Mode: d.Get("mode").(string), + TeamID: d.Get("team_id").(string), + DisplayName: d.Get("display_name").(string), }, Additional: make(map[string]interface{}), } @@ -275,6 +278,12 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return fmt.Errorf("failed to %s model: %w", map[bool]string{true: "update", false: "create"}[isUpdate], err) } + if isUpdate && d.HasChange("display_name") { + if err := patchModelDisplayName(client, modelID, d.Get("display_name").(string)); err != nil { + return fmt.Errorf("failed to update model display_name: %w", err) + } + } + d.SetId(modelID) log.Printf("[INFO] Model created with ID %s. Starting retry mechanism to read the model...", modelID) @@ -282,6 +291,19 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return retryModelRead(d, m, 5) } +// /model/update only merges litellm_params, so model_info changes go through the PATCH endpoint. +func patchModelDisplayName(client *Client, modelID, displayName string) error { + resp, err := MakeRequest(client, "PATCH", fmt.Sprintf(endpointModelPatch, url.PathEscape(modelID)), ModelInfoPatch{ + ModelInfo: ModelInfoPatchFields{ID: modelID, DisplayName: displayName}, + }) + if err != nil { + return err + } + defer resp.Body.Close() + _, err = handleAPIResponse(resp, nil, client) + return err +} + func resourceLiteLLMModelCreate(d *schema.ResourceData, m interface{}) error { return createOrUpdateModel(d, m, false) } @@ -327,6 +349,7 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string))) d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string))) d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string))) + d.Set("display_name", modelResp.ModelInfo.DisplayName) // Preserve credential name from state since it might not be returned by API d.Set("litellm_credential_name", d.Get("litellm_credential_name").(string)) diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go new file mode 100644 index 00000000000..0be3c39a3c5 --- /dev/null +++ b/terraform/provider/litellm/resource_model_test.go @@ -0,0 +1,244 @@ +package litellm + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +func modelInfoBody(displayName string) string { + modelInfo := map[string]interface{}{ + "id": "model-123", + "db_model": true, + "base_model": "claude-sonnet-4-5", + "tier": "free", + "mode": "chat", + } + if displayName != "" { + modelInfo["display_name"] = displayName + } + body, _ := json.Marshal(map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "litellm_params": map[string]interface{}{"model": "anthropic/claude-sonnet-4-5", "custom_llm_provider": "anthropic"}, + "model_info": modelInfo, + }) + return string(body) +} + +func modelInfoDataEnvelope(displayName string) string { + return `{"data": [` + modelInfoBody(displayName) + `]}` +} + +func TestResourceLiteLLMModelCreateSendsDisplayName(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + case "/model/info": + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "model_api_key": "sk-ant-test", + "mode": "chat", + "display_name": "Claude Sonnet 4.5", + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + modelInfo, ok := createPayload["model_info"].(map[string]interface{}) + if !ok { + t.Fatalf("expected model_info object in create payload, got %v", createPayload["model_info"]) + } + if modelInfo["display_name"] != "Claude Sonnet 4.5" { + t.Errorf("expected model_info.display_name 'Claude Sonnet 4.5', got %v", modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != "Claude Sonnet 4.5" { + t.Errorf("expected state display_name 'Claude Sonnet 4.5', got %q", got) + } +} + +func TestResourceLiteLLMModelCreateOmitsUnsetDisplayName(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "model_api_key": "sk-ant-test", + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + modelInfo := createPayload["model_info"].(map[string]interface{}) + if _, present := modelInfo["display_name"]; present { + t.Errorf("expected display_name to be omitted from model_info when unset, got %v", modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != "" { + t.Errorf("expected empty state display_name, got %q", got) + } +} + +func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { + cases := map[string]struct { + serverBody string + want string + }{ + "server value wins inside data envelope": {serverBody: modelInfoDataEnvelope("Renamed In Admin UI"), want: "Renamed In Admin UI"}, + "server value wins unwrapped": {serverBody: modelInfoBody("Renamed In Admin UI"), want: "Renamed In Admin UI"}, + "external removal clears state": {serverBody: modelInfoDataEnvelope(""), want: ""}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/model/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Write([]byte(tc.serverBody)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "display_name": "Claude Sonnet 4.5", + }) + d.SetId("model-123") + + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if got := d.Get("display_name").(string); got != tc.want { + t.Errorf("expected display_name %q, got %q", tc.want, got) + } + }) + } +} + +func updateResourceData(t *testing.T, oldDisplayName, newDisplayName string) *schema.ResourceData { + t.Helper() + res := resourceLiteLLMModel() + attrs := map[string]string{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + } + if oldDisplayName != "" { + attrs["display_name"] = oldDisplayName + } + state := &terraform.InstanceState{ID: "model-123", Attributes: attrs} + diff, err := res.Diff(context.Background(), state, &terraform.ResourceConfig{Config: map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "display_name": newDisplayName, + }}, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(state, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + return d +} + +func TestResourceLiteLLMModelUpdatePatchesDisplayName(t *testing.T) { + cases := map[string]struct { + newName string + }{ + "changed name is patched": {newName: "Claude Sonnet 4.5 v2"}, + "cleared name is patched": {newName: ""}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + var patchPayload map[string]interface{} + var patchPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/model/update": + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + case r.Method == http.MethodPatch: + patchPath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&patchPayload); err != nil { + t.Errorf("failed to decode patch payload: %v", err) + } + w.Write([]byte(modelInfoBody(tc.newName))) + case r.URL.Path == "/model/info": + w.Write([]byte(modelInfoDataEnvelope(tc.newName))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := updateResourceData(t, "Claude Sonnet 4.5", tc.newName) + if err := resourceLiteLLMModelUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + if patchPath != "/model/model-123/update" { + t.Fatalf("expected PATCH /model/model-123/update, got %q", patchPath) + } + modelInfo := patchPayload["model_info"].(map[string]interface{}) + if modelInfo["display_name"] != tc.newName { + t.Errorf("expected patched display_name %q, got %v", tc.newName, modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != tc.newName { + t.Errorf("expected state display_name %q, got %q", tc.newName, got) + } + }) + } +} + +func TestResourceLiteLLMModelUpdateSkipsPatchWhenDisplayNameUnchanged(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPatch { + t.Errorf("unexpected PATCH %s", r.URL.Path) + } + w.Write([]byte(modelInfoDataEnvelope("Claude Sonnet 4.5"))) + })) + defer srv.Close() + + d := updateResourceData(t, "Claude Sonnet 4.5", "Claude Sonnet 4.5") + if err := resourceLiteLLMModelUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 8bcf7dc4fe3..a8784b8a6a9 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -25,6 +25,16 @@ type ModelResponse struct { Additional map[string]interface{} `json:"additional"` } +// ModelInfoPatch is the body for PATCH /model/{id}/update; display_name is sent even when empty so it can be cleared. +type ModelInfoPatch struct { + ModelInfo ModelInfoPatchFields `json:"model_info"` +} + +type ModelInfoPatchFields struct { + ID string `json:"id"` + DisplayName string `json:"display_name"` +} + // ModelRequest represents a request to create or update a model. type ModelRequest struct { ModelName string `json:"model_name"` @@ -108,12 +118,13 @@ type LiteLLMParams struct { // ModelInfo represents information about a model. type ModelInfo struct { - ID string `json:"id"` - DBModel bool `json:"db_model"` - BaseModel string `json:"base_model"` - Tier string `json:"tier"` - Mode string `json:"mode"` - TeamID string `json:"team_id,omitempty"` + ID string `json:"id"` + DBModel bool `json:"db_model"` + BaseModel string `json:"base_model"` + Tier string `json:"tier"` + Mode string `json:"mode"` + TeamID string `json:"team_id,omitempty"` + DisplayName string `json:"display_name,omitempty"` } // Key represents a LiteLLM API key. diff --git a/terraform/provider/litellm/utils.go b/terraform/provider/litellm/utils.go index f8f66afba3c..ce1dae55f59 100644 --- a/terraform/provider/litellm/utils.go +++ b/terraform/provider/litellm/utils.go @@ -55,6 +55,13 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client) resp.Status, client.redactSensitiveData(string(bodyBytes)), client.redactSensitiveData(string(reqBodyBytes))) } + var envelope struct { + Data []json.RawMessage `json:"data"` + } + if err := json.Unmarshal(bodyBytes, &envelope); err == nil && len(envelope.Data) > 0 { + bodyBytes = envelope.Data[0] + } + var modelResp ModelResponse if err := json.Unmarshal(bodyBytes, &modelResp); err != nil { return nil, fmt.Errorf("failed to parse response: %v", err) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 7aacf7ceab9..e4574031d86 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -97,7 +97,6 @@ GET /guardrails/{guardrail_id} GET /prompts/{prompt_id} GET /prompts/{prompt_id}/versions PATCH /guardrails/{guardrail_id} -PATCH /model/{model_id}/update PATCH /prompts/{prompt_id} PATCH /team/{team_id} POST /team/model/add From de8aeff6c6d2f360b4ade529ecb90159c896bf1e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:07:10 -0500 Subject: [PATCH 021/154] feat(proxy_cli): add --validate_config dry-run flag (#41705) * feat(proxy_cli): add --validate_config dry-run flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(tests): format --validate_config CliRunner calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy_cli): run --validate_config before the ollama auto-start Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy_cli): restore file and add ollama validate_config regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin --- litellm/proxy/proxy_cli.py | 29 ++++++ tests/test_litellm/proxy/test_proxy_cli.py | 103 +++++++++++++++++++++ 2 files changed, 132 insertions(+) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ac69f2e3894..27c03d2d5d7 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -210,6 +210,25 @@ class ProxyInitializationHelpers: response: Final = httpx.get(url=f"http://{host}:{port}/health") print(json.dumps(response.json(), indent=4)) + @staticmethod + def _run_config_validation(config: str | None) -> None: + if config is None: + raise click.UsageError("--validate_config requires --config ") + import asyncio + + from litellm.proxy.proxy_server import ProxyConfig + + async def _load() -> int: + _, model_list, _ = await ProxyConfig().load_config(router=None, config_file_path=config) + return len(model_list) + + try: + model_count: Final = asyncio.run(_load()) + except Exception as error: + click.echo(f"LiteLLM: config validation failed: {error}", err=True) + raise click.exceptions.Exit(1) from error + click.echo(f"LiteLLM: config OK ({model_count} models)") + @staticmethod def _run_test_chat_completion( host: str, @@ -887,6 +906,12 @@ class ProxyInitializationHelpers: default=False, help="Skip starting the server after setup (useful for migrations only)", ) +@click.option( + "--validate_config", + is_flag=True, + default=False, + help="Load and validate the config file (including mcp_servers) without starting the server, then exit. Exit code 1 on any config error.", +) @click.option( "--keepalive_timeout", default=None, @@ -1027,6 +1052,7 @@ def run_server( log_config, use_prisma_db_push: bool, skip_server_startup, + validate_config: bool, keepalive_timeout, timeout_worker_healthcheck, max_requests_before_restart, @@ -1069,6 +1095,9 @@ def run_server( if version is True: ProxyInitializationHelpers._echo_litellm_version() return + if validate_config is True: + ProxyInitializationHelpers._run_config_validation(config) + return if model and "ollama" in model and api_base is None: ProxyInitializationHelpers._run_ollama_serve() if health is True: diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index b84efda3308..a275dd62400 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -3224,3 +3224,106 @@ class TestLibpqSslParamTranslation: assert query["sslmode"] == ["require"] assert query["sslcert"] == ["/certs/rds-bundle.pem"] assert query["sslaccept"] == ["strict"] + + +@pytest.mark.xdist_group("proxy_cli") +class TestValidateConfigFlag: + def test_validate_config_valid_config_exits_zero(self, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ] + } + ) + ) + + result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"]) + + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert "config OK" in result.output + + def test_validate_config_invalid_mcp_server_exits_one(self, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "mcp_servers": { + "zapier": { + "url": "https://example.com/mcp", + "transport": "http", + "per_server_oauth_discovery": "yes", + } + } + } + ) + ) + + result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"]) + + assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}" + assert "per_server_oauth_discovery must be a boolean" in result.output + + def test_validate_config_without_config_is_usage_error(self, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + result = CliRunner().invoke(run_server, ["--validate_config"]) + + assert result.exit_code != 0 + assert "--validate_config requires --config" in result.output + + @patch("subprocess.Popen") + def test_validate_config_with_ollama_model_does_not_start_ollama(self, mock_popen, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ] + } + ) + ) + + result = CliRunner().invoke( + run_server, + ["--config", str(config_path), "--model", "ollama/llama3", "--validate_config"], + ) + + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert "config OK" in result.output + mock_popen.assert_not_called() From 7faeb15ff3dfa5e227e32fb69b3a2b6945b3ad72 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:17:31 -0700 Subject: [PATCH 022/154] fix(e2e): skip unpublished npm versions in the Claude Code PR-gate resolver (#43053) npm keeps an unpublished version's timestamp in the packument's time map but drops it from versions, so the resolver could hand npm install a version it refuses with ETARGET. Only versions still present in versions are candidates now. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../_pr_gate_unit_tests/__init__.py | 0 .../test_pr_gate_version_resolver.py | 60 +++++++++++++++++++ .../claude_code/pr_gate_version_resolver.py | 7 ++- 3 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py b/tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py b/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py new file mode 100644 index 00000000000..0ed7e2bf083 --- /dev/null +++ b/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py @@ -0,0 +1,60 @@ +"""Unit tests for the Claude Code PR-gate version resolver. + +Markerless harness tests: they feed the resolver a hand-built packument and a +fixed clock, so they run without a proxy, never reach the npm registry, and +carry no `e2e` marker. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Final, Mapping + +import pytest + +from claude_code.pr_gate_version_resolver import NoEligibleVersionError, resolve_pr_gate_version + +NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) +INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) + + +def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> dict[str, object]: + return { + "name": "@anthropic-ai/claude-code", + "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, + "versions": {version: {"version": version} for version in times if version not in unpublished}, + } + + +def test_skips_a_version_npm_has_unpublished() -> None: + metadata: Final = _packument( + { + "2.1.87": "2026-03-28T20:00:00.000Z", + "2.1.88": "2026-03-30T22:36:48.424Z", + "2.1.89": "2026-03-31T23:32:40.000Z", + }, + unpublished=frozenset({"2.1.88"}), + ) + assert resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" + + +def test_raises_when_the_only_old_enough_version_is_unpublished() -> None: + metadata: Final = _packument( + {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, + unpublished=frozenset({"2.1.88"}), + ) + with pytest.raises(NoEligibleVersionError): + resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) + + +def test_picks_the_newest_published_version_at_least_min_age_old() -> None: + metadata: Final = _packument( + { + "2.1.118": "2026-04-15T10:00:00.000Z", + "2.1.119": "2026-04-21T10:00:00.000Z", + "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", + "2.1.120": "2026-04-23T10:00:00.000Z", + "2.1.121": "2026-04-25T11:00:00.000Z", + } + ) + assert resolve_pr_gate_version(metadata=metadata, as_of=NOW) == "2.1.119" diff --git a/tests/e2e/claude_code/pr_gate_version_resolver.py b/tests/e2e/claude_code/pr_gate_version_resolver.py index 82e12a2bf15..756dafb7d73 100644 --- a/tests/e2e/claude_code/pr_gate_version_resolver.py +++ b/tests/e2e/claude_code/pr_gate_version_resolver.py @@ -80,7 +80,9 @@ def resolve_pr_gate_version( "Newest" means newest by **publish time**, not semver string order — if a patch lands on an older major after a newer release, the - patched line is the eligible one. + patched line is the eligible one. A version npm has unpublished keeps + its ``time`` entry but drops out of ``versions``, so only versions + still present in ``versions`` are candidates. Args: metadata: Pre-fetched npm packument (skips the HTTP call). Useful @@ -101,6 +103,7 @@ def resolve_pr_gate_version( metadata = fetch(package_name) times = metadata.get("time") or {} + versions = metadata.get("versions") or {} if as_of is None: as_of = datetime.now(timezone.utc) cutoff = as_of - min_age @@ -109,6 +112,8 @@ def resolve_pr_gate_version( for version, raw_ts in times.items(): if version in _TIME_META_KEYS: continue + if version not in versions: + continue if not isinstance(raw_ts, str): continue if "-" in version: From 040b37fa49b99be78b997f76ad46bd6f10c07f27 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:18:55 -0700 Subject: [PATCH 023/154] chore(cost-map): move azure gpt-realtime-2.1-mini deprecation date to the later Models API date (#43058) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 680b6cceab6..9bec83f6b08 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6092,7 +6092,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 680b6cceab6..9bec83f6b08 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6092,7 +6092,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, From 4aa3ff47fe1fd525c54edd9adc4e206cf439baea Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:23:05 -0700 Subject: [PATCH 024/154] docs(github): require UI before/after screenshots and intentional UX change note in PR template (#43021) * docs(github): require UI before/after screenshots and intentional UX change note in PR template Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(github): move intentional change note into TLDR rules and dedupe screenshots Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(github): refresh user flow screenshots with new commits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/pull_request_template.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 4b3878bed11..1fe0c602036 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -6,7 +6,8 @@ ## TLDR - + Problem this solves: @@ -28,7 +29,8 @@ How it solves it: No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong Keep the two lists step-for-step identical until they diverge, so the changed step is obvious If the bug had a security or authorization consequence, end each list with what another user could or could no longer do - Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision + Regenerate this section, screenshots included, whenever new commits change the PR's behavior, so it never describes an older revision + If the PR changes what an Admin UI page shows, embed a before and an after screenshot of that page right after its list, taken at the same URL on the same data, with the rows, fields, or controls that changed boxed in red so a reader spots the difference without reading the steps. These are the UI screenshots for Screenshots / Proof of Fix too: embed them once here and have that section's Before and After steps point back to them instead of repeating the images Example: From bf0187072bb360153400a17895e65dc10f4a4110 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:49:59 -0700 Subject: [PATCH 025/154] ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests (#42902) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/unit_selection.sh | 63 +++++++++++++++++++ .circleci/tests.yml | 33 +++++++++- .github/scripts/assert_ci_coverage.py | 12 +++- .github/workflows/_test-unit-base.yml | 33 ++++++++-- .github/workflows/test-unit.yml | 22 ++++--- Makefile | 2 +- tests/test_litellm/test_assert_ci_coverage.py | 2 +- .../proxy => unit/caching}/__init__.py | 0 .../caching}/test_cache_preset_key.py | 0 .../caching}/test_caching_handler.py | 0 .../test_responses_stream_cache_keys.py | 0 .../caching}/test_unit_test_caching.py | 0 tests/unit/conftest.py | 2 +- tests/{ => unit}/enterprise/conftest.py | 0 .../send_emails/__init__.py | 0 .../send_emails/test_base_email.py | 0 .../send_emails/test_endpoints.py | 0 .../send_emails/test_resend_email.py | 0 .../send_emails/test_sendgrid_email.py | 0 .../test_prometheus_logging_callbacks.py | 0 .../unit/enterprise/integrations/__init__.py | 0 .../integrations/test_custom_guardrail.py | 0 .../integrations/test_prometheus.py | 0 .../test_prometheus_unit_tests.py | 0 tests/unit/enterprise/proxy/__init__.py | 0 tests/unit/enterprise/proxy/auth/__init__.py | 0 .../proxy/auth/test_route_checks.py | 0 .../proxy/auth/test_user_api_key_auth.py | 0 .../enterprise/proxy/guardrails/__init__.py | 0 .../enterprise}/proxy/guardrails/conftest.py | 0 .../test_apply_guardrail_endpoint.py | 0 .../test_bedrock_apply_guardrail.py | 0 tests/unit/enterprise/proxy/hooks/__init__.py | 0 .../proxy/hooks/test_managed_files.py | 0 .../proxy/management_endpoints/__init__.py | 0 .../test_internal_user_endpoints.py | 0 .../test_project_endpoints_prisma.py | 0 .../test_afile_retrieve_returns_unified_id.py | 0 .../proxy/test_audit_logging_endpoints.py | 0 .../test_batch_retrieve_input_file_id.py | 0 ...trieve_registers_missing_output_file_id.py | 0 ..._retrieve_returns_unified_input_file_id.py | 0 ..._batch_update_db_managed_output_file_id.py | 0 .../test_deleted_file_returns_403_not_404.py | 0 .../proxy/test_enterprise_routes.py | 0 .../proxy/test_file_deletion_blocking.py | 0 .../proxy/test_managed_files_access_check.py | 0 .../proxy/test_managed_files_hook.py | 0 tests/unit/gateway/__init__.py | 0 .../gateway}/test_launch.py | 0 tests/unit/litellm_proxy_extras/__init__.py | 0 .../test_litellm_proxy_extras_logging.py | 0 .../test_litellm_proxy_extras_utils.py | 6 +- 53 files changed, 152 insertions(+), 23 deletions(-) create mode 100755 .circleci/scripts/unit_selection.sh rename tests/{test_litellm/enterprise/proxy => unit/caching}/__init__.py (100%) rename tests/{local_testing => unit/caching}/test_cache_preset_key.py (100%) rename tests/{local_testing => unit/caching}/test_caching_handler.py (100%) rename tests/{local_testing => unit/caching}/test_responses_stream_cache_keys.py (100%) rename tests/{local_testing => unit/caching}/test_unit_test_caching.py (100%) rename tests/{ => unit}/enterprise/conftest.py (100%) create mode 100644 tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_base_email.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_endpoints.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_resend_email.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/enterprise_callbacks/test_prometheus_logging_callbacks.py (100%) create mode 100644 tests/unit/enterprise/integrations/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_custom_guardrail.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_prometheus.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_prometheus_unit_tests.py (100%) create mode 100644 tests/unit/enterprise/proxy/__init__.py create mode 100644 tests/unit/enterprise/proxy/auth/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/auth/test_route_checks.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/auth/test_user_api_key_auth.py (100%) create mode 100644 tests/unit/enterprise/proxy/guardrails/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/conftest.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/test_apply_guardrail_endpoint.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/test_bedrock_apply_guardrail.py (100%) create mode 100644 tests/unit/enterprise/proxy/hooks/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/hooks/test_managed_files.py (100%) create mode 100644 tests/unit/enterprise/proxy/management_endpoints/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/management_endpoints/test_internal_user_endpoints.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/management_endpoints/test_project_endpoints_prisma.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_afile_retrieve_returns_unified_id.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/test_audit_logging_endpoints.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_input_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_update_db_managed_output_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_deleted_file_returns_403_not_404.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_enterprise_routes.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_file_deletion_blocking.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_managed_files_access_check.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_managed_files_hook.py (100%) create mode 100644 tests/unit/gateway/__init__.py rename tests/{test_gateway => unit/gateway}/test_launch.py (100%) create mode 100644 tests/unit/litellm_proxy_extras/__init__.py rename tests/{litellm-proxy-extras => unit/litellm_proxy_extras}/test_litellm_proxy_extras_logging.py (100%) rename tests/{litellm-proxy-extras => unit/litellm_proxy_extras}/test_litellm_proxy_extras_utils.py (99%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh new file mode 100755 index 00000000000..8d2b8a42691 --- /dev/null +++ b/.circleci/scripts/unit_selection.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +set -euo pipefail + +flag="${1:?usage: unit_selection.sh }" + +legacy_flags=( + caching-local + enterprise-package + enterprise-routing + proxy-extras + proxy-infra +) + +legacy_paths() { + case "$1" in + caching-local) echo tests/unit/caching ;; + enterprise-package) + echo tests/unit/enterprise/integrations + echo tests/unit/enterprise/proxy/auth + echo tests/unit/enterprise/proxy/guardrails + echo tests/unit/enterprise/proxy/hooks + echo tests/unit/enterprise/proxy/management_endpoints + echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py + echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; + enterprise-routing) + echo tests/unit/enterprise/enterprise_callbacks/send_emails + echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py + echo tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py + echo tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py + echo tests/unit/enterprise/proxy/test_enterprise_routes.py + echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py + echo tests/unit/enterprise/proxy/test_managed_files_access_check.py + echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + proxy-extras) echo tests/unit/litellm_proxy_extras ;; + proxy-infra) echo tests/unit/gateway ;; + *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; + esac +} + +expand() { + while read -r path; do + if [ -d "$path" ]; then + find "$path" -name 'test_*.py' + elif [ -f "$path" ]; then + echo "$path" + else + echo "unit_selection.sh: $path does not exist" >&2 + exit 1 + fi + done +} + +if [ "$flag" = unit ]; then + comm -23 \ + <(find tests/unit -name 'test_*.py' | sort) \ + <(for legacy in "${legacy_flags[@]}"; do legacy_paths "$legacy"; done | expand | sort) + exit 0 +fi + +legacy_paths "$flag" | expand | sort diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 1afb935453f..38fe44bb625 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -171,6 +171,12 @@ jobs: shards: type: integer default: 6 + workers: + type: integer + default: 4 + dist: + type: string + default: loadscope base_ref: type: string default: "" @@ -199,17 +205,19 @@ jobs: no_output_timeout: 20m command: | mkdir -p test-results/<< parameters.flag >> - selection="$(find tests/unit -name 'test_*.py' | sort)" || { echo "test selection failed for << parameters.flag >>"; exit 1; } - [ -n "${selection}" ] || { echo "test selection produced no files for << parameters.flag >>"; exit 1; } + selection="$(bash .circleci/scripts/unit_selection.sh << parameters.flag >>)" || { echo "unit_selection.sh failed for << parameters.flag >>"; exit 1; } + [ -n "${selection}" ] || { echo "unit_selection.sh produced no files for << parameters.flag >>"; exit 1; } shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; } [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; } mapfile -t files < <(printf '%s\n' "${shard}") + xdist_args=() + if [ "<< parameters.workers >>" -gt 0 ]; then xdist_args=(-n << parameters.workers >> --dist=<< parameters.dist >>); fi rerun_args=(-p no:rerunfailures) if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") set +e env -i "${test_env[@]}" \ - uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml + uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml status=$? set -e if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi @@ -293,6 +301,25 @@ workflows: - unit: base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-<< matrix.flag >> + shards: 1 + workers: 2 + reruns: 2 + matrix: + parameters: + flag: [caching-local, proxy-extras, enterprise-routing] + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-<< matrix.flag >> + shards: 1 + reruns: 2 + matrix: + parameters: + flag: [enterprise-package, proxy-infra] + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - documentation - integration: name: integration-<< matrix.suite >> diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index 2e008fe7ade..d8246225a3b 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -120,6 +120,13 @@ def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: ) +def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]: + script: Final = repo_root / ".circleci/scripts/unit_selection.sh" + if not script.is_file(): + return frozenset() + return frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(_uncommented(script.read_text()))) + + def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: return frozenset( match.group(0) @@ -611,7 +618,10 @@ def main() -> int: scalars = _all_scalars() integration_paths, ownership_findings = _integration_ownership() - test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | integration_paths) + ownership_findings + test_findings = ( + _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | _unit_selection_tokens() | integration_paths) + + ownership_findings + ) dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars)) stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles()) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 8faddd11df8..ef1dc53b4a6 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -13,6 +13,15 @@ on: have its path existence-checked like any other token. required: true type: string + fork-flag: + description: >- + Codecov flag of the `.circleci/tests.yml` job that now owns part of + this shard. CircleCI does not run on pull requests from forks, so on + those events this shard also runs the files + `.circleci/scripts/unit_selection.sh` lists for the flag. + required: false + type: string + default: "" workers: description: "Number of pytest-xdist workers" required: false @@ -92,6 +101,7 @@ jobs: pull-requests: read outputs: decision: ${{ steps.changes.outputs.decision }} + has-coverage: ${{ steps.tests.outputs.has-coverage }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -160,10 +170,13 @@ jobs: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - name: Run tests + id: tests if: steps.changes.outputs.decision != 'skip' timeout-minutes: ${{ inputs.timeout-minutes }} env: TEST_PATH: ${{ inputs.test-path }} + FORK_FLAG: ${{ inputs.fork-flag }} + IS_FORK: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository }} MAX_FAILURES: ${{ inputs.max-failures }} WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} @@ -171,9 +184,18 @@ jobs: DIST: ${{ inputs.dist }} COVERAGE_CORE: sysmon run: | + echo "has-coverage=false" >> "$GITHUB_OUTPUT" + selection="${TEST_PATH}" + if [ "${IS_FORK}" = "true" ] && [ -n "${FORK_FLAG}" ]; then + selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${FORK_FLAG}" | tr '\n' ' ')" + fi + if [ -z "${selection// /}" ]; then + echo "shard selection is empty on this event (CircleCI flag ${FORK_FLAG:-none} owns it); nothing to run" + exit 0 + fi pytest_args=() existing_paths=0 - for token in ${TEST_PATH:?}; do + for token in ${selection}; do case "${token}" in -*) pytest_args+=("${token}") ;; *) @@ -187,7 +209,7 @@ jobs: esac done if [ "${existing_paths}" -eq 0 ]; then - echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run" + echo "No path in the selection exists (${selection}); nothing to run" exit 0 fi xdist_args=() @@ -209,8 +231,11 @@ jobs: --cov-config=pyproject.toml status=$? set -e + if [ -f coverage.xml ]; then + echo "has-coverage=true" >> "$GITHUB_OUTPUT" + fi if [ "$status" -eq 5 ]; then - echo "pytest collected no tests from ${TEST_PATH}; passing" + echo "pytest collected no tests from ${selection}; passing" exit 0 fi exit "$status" @@ -226,7 +251,7 @@ jobs: upload-coverage: name: Upload coverage to Codecov needs: run - if: always() && needs.run.outputs.decision != 'skip' + if: always() && needs.run.outputs.decision != 'skip' && needs.run.outputs.has-coverage == 'true' runs-on: ubuntu-latest permissions: contents: read diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 686bbc89467..94de6040038 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -35,6 +35,10 @@ concurrency: # already a matrix and carries a shard-coverage guard that reads that file by # name. Folding it in here is a follow-up, together with generalising that guard # into assert_ci_coverage.py. +# +# `fork-flag` names the `.circleci/tests.yml` job that now runs part of the +# shard under the same Codecov flag. CircleCI does not build pull requests from +# forks, so the shard still runs those files there and skips them elsewhere. jobs: unit: name: ${{ matrix.shard }} @@ -65,10 +69,10 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing test-path: >- - tests/test_litellm/enterprise tests/test_litellm/google_genai tests/test_litellm/router_utils tests/test_litellm/router_strategy + fork-flag: enterprise-routing workers: 2 reruns: 2 timeout-minutes: 20 @@ -200,7 +204,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py - tests/test_gateway + fork-flag: proxy-infra workers: 4 reruns: 2 timeout-minutes: 20 @@ -208,11 +212,8 @@ jobs: - shard: caching-local artifact-name: caching-local - test-path: >- - tests/local_testing/test_cache_preset_key.py - tests/local_testing/test_caching_handler.py - tests/local_testing/test_responses_stream_cache_keys.py - tests/local_testing/test_unit_test_caching.py + test-path: "" + fork-flag: caching-local workers: 2 reruns: 2 timeout-minutes: 20 @@ -220,7 +221,8 @@ jobs: - shard: proxy-extras artifact-name: proxy-extras - test-path: "tests/litellm-proxy-extras" + test-path: "" + fork-flag: proxy-extras workers: 2 reruns: 2 timeout-minutes: 20 @@ -228,7 +230,8 @@ jobs: - shard: enterprise-package artifact-name: enterprise-package - test-path: "tests/enterprise" + test-path: "" + fork-flag: enterprise-package workers: 4 reruns: 2 timeout-minutes: 20 @@ -247,6 +250,7 @@ jobs: uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} + fork-flag: ${{ matrix.fork-flag || '' }} workers: ${{ matrix.workers }} reruns: ${{ matrix.reruns }} timeout-minutes: ${{ matrix.timeout-minutes }} diff --git a/Makefile b/Makefile index ab7fab6aa99..6263b646c17 100644 --- a/Makefile +++ b/Makefile @@ -332,7 +332,7 @@ test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/tests/test_litellm/test_assert_ci_coverage.py b/tests/test_litellm/test_assert_ci_coverage.py index 983707db606..8524a905745 100644 --- a/tests/test_litellm/test_assert_ci_coverage.py +++ b/tests/test_litellm/test_assert_ci_coverage.py @@ -340,7 +340,7 @@ def test_a_dockerfile_directory_entry_is_stale_because_only_an_exact_path_exempt def test_a_workflow_that_names_a_file_clears_it_from_the_slice_check(): named = coverage._workflow_named_tokens() assert named, "the workflows must name some test paths or the check proves nothing" - assert any(coverage._token_covers(token, "tests/local_testing/test_caching_handler.py") for token in named) + assert any(coverage._token_covers(token, "tests/proxy_unit_tests/test_proxy_custom_logger.py") for token in named) def test_the_slice_check_credits_only_workflows_never_the_circleci_config(): diff --git a/tests/test_litellm/enterprise/proxy/__init__.py b/tests/unit/caching/__init__.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/__init__.py rename to tests/unit/caching/__init__.py diff --git a/tests/local_testing/test_cache_preset_key.py b/tests/unit/caching/test_cache_preset_key.py similarity index 100% rename from tests/local_testing/test_cache_preset_key.py rename to tests/unit/caching/test_cache_preset_key.py diff --git a/tests/local_testing/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py similarity index 100% rename from tests/local_testing/test_caching_handler.py rename to tests/unit/caching/test_caching_handler.py diff --git a/tests/local_testing/test_responses_stream_cache_keys.py b/tests/unit/caching/test_responses_stream_cache_keys.py similarity index 100% rename from tests/local_testing/test_responses_stream_cache_keys.py rename to tests/unit/caching/test_responses_stream_cache_keys.py diff --git a/tests/local_testing/test_unit_test_caching.py b/tests/unit/caching/test_unit_test_caching.py similarity index 100% rename from tests/local_testing/test_unit_test_caching.py rename to tests/unit/caching/test_unit_test_caching.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index b3bb19a8b8a..202ecb80d7b 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -11,7 +11,7 @@ import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency -LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( "AZURE_AD_TOKEN", "AZURE_TENANT_ID", diff --git a/tests/enterprise/conftest.py b/tests/unit/enterprise/conftest.py similarity index 100% rename from tests/enterprise/conftest.py rename to tests/unit/enterprise/conftest.py diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_endpoints.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_endpoints.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_resend_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_resend_email.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py similarity index 100% rename from tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py rename to tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py diff --git a/tests/unit/enterprise/integrations/__init__.py b/tests/unit/enterprise/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py b/tests/unit/enterprise/integrations/test_custom_guardrail.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py rename to tests/unit/enterprise/integrations/test_custom_guardrail.py diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py b/tests/unit/enterprise/integrations/test_prometheus.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_prometheus.py rename to tests/unit/enterprise/integrations/test_prometheus.py diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py rename to tests/unit/enterprise/integrations/test_prometheus_unit_tests.py diff --git a/tests/unit/enterprise/proxy/__init__.py b/tests/unit/enterprise/proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/proxy/auth/__init__.py b/tests/unit/enterprise/proxy/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py b/tests/unit/enterprise/proxy/auth/test_route_checks.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py rename to tests/unit/enterprise/proxy/auth/test_route_checks.py diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/unit/enterprise/proxy/auth/test_user_api_key_auth.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py rename to tests/unit/enterprise/proxy/auth/test_user_api_key_auth.py diff --git a/tests/unit/enterprise/proxy/guardrails/__init__.py b/tests/unit/enterprise/proxy/guardrails/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py b/tests/unit/enterprise/proxy/guardrails/conftest.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py rename to tests/unit/enterprise/proxy/guardrails/conftest.py diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py rename to tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/unit/enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py rename to tests/unit/enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py diff --git a/tests/unit/enterprise/proxy/hooks/__init__.py b/tests/unit/enterprise/proxy/hooks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py rename to tests/unit/enterprise/proxy/hooks/test_managed_files.py diff --git a/tests/unit/enterprise/proxy/management_endpoints/__init__.py b/tests/unit/enterprise/proxy/management_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/enterprise/proxy/management_endpoints/test_internal_user_endpoints.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py rename to tests/unit/enterprise/proxy/management_endpoints/test_internal_user_endpoints.py diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py rename to tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py diff --git a/tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py b/tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py rename to tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/unit/enterprise/proxy/test_audit_logging_endpoints.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py rename to tests/unit/enterprise/proxy/test_audit_logging_endpoints.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py rename to tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py b/tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py rename to tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py diff --git a/tests/test_litellm/enterprise/proxy/test_enterprise_routes.py b/tests/unit/enterprise/proxy/test_enterprise_routes.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_enterprise_routes.py rename to tests/unit/enterprise/proxy/test_enterprise_routes.py diff --git a/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py rename to tests/unit/enterprise/proxy/test_file_deletion_blocking.py diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/unit/enterprise/proxy/test_managed_files_access_check.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py rename to tests/unit/enterprise/proxy/test_managed_files_access_check.py diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_managed_files_hook.py rename to tests/unit/enterprise/proxy/test_managed_files_hook.py diff --git a/tests/unit/gateway/__init__.py b/tests/unit/gateway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_gateway/test_launch.py b/tests/unit/gateway/test_launch.py similarity index 100% rename from tests/test_gateway/test_launch.py rename to tests/unit/gateway/test_launch.py diff --git a/tests/unit/litellm_proxy_extras/__init__.py b/tests/unit/litellm_proxy_extras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_logging.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_logging.py similarity index 100% rename from tests/litellm-proxy-extras/test_litellm_proxy_extras_logging.py rename to tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_logging.py diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py similarity index 99% rename from tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py rename to tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index bb329264a11..755c7617701 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -9,7 +9,7 @@ import pytest sys.path.insert( 0, os.path.abspath( - os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") + os.path.join(os.path.dirname(__file__), "../../../litellm-proxy-extras") ), ) @@ -23,7 +23,7 @@ from litellm_proxy_extras.utils import ( _MIGRATIONS_DIR = os.path.abspath( os.path.join( os.path.dirname(__file__), - "../../litellm-proxy-extras/litellm_proxy_extras/migrations", + "../../../litellm-proxy-extras/litellm_proxy_extras/migrations", ) ) @@ -999,7 +999,7 @@ class TestJWTKeyMappingCascade: schema_paths = glob.glob( os.path.abspath( os.path.join( - os.path.dirname(__file__), "../../**/schema.prisma" + os.path.dirname(__file__), "../../../**/schema.prisma" ) ), recursive=True, From 7b25a151bd72dbed46137a89799b298d6da4d87e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:52:44 -0700 Subject: [PATCH 026/154] feat(proxy): let callbacks filter the model listing routes per caller (#43027) * feat(proxy): let callbacks filter the model listing routes per caller * fix(proxy): offer every listed name to the listing callback, agent groups and deployment lookups included * fix(proxy): hide aliases of a team model by its public name and offer /model/info lookups the listed name * fix(proxy): map a malformed model listing filter return to the proxy error contract and document legacy team names --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_logger.py | 18 + litellm/proxy/proxy_server.py | 101 ++++- litellm/proxy/utils.py | 62 ++- .../proxy_server/test_routes_model_info.py | 13 +- .../proxy/test_model_list_callback_filter.py | 425 ++++++++++++++++++ 5 files changed, 596 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/proxy/test_model_list_callback_filter.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 326abd5c6a3..d4162369a35 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -421,6 +421,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm pass + async def async_filter_listed_models( + self, + user_api_key_dict: UserAPIKeyAuth, + model_names: Sequence[str], + ) -> Sequence[str]: + """Runs on the model listing routes (`/v1/models`, `/v1/models/{id}`, `/model/info`, + `/model_group/info`) with the public model names the route would otherwise return, so a + lookup of one model may offer just that name: decide per name, never by position in the + sequence. Return the names to keep as a sequence of strings; a name left out disappears + from every listing, any alias of it offered in the same call goes with it, and + `/v1/models/{id}` answers 404 for it, exactly as for a model that does not exist. Names + outside `model_names` are ignored, so a callback can only narrow the listing, never widen + it. Under `use_team_public_model_name: false`, `/v1/models` and `/model_group/info` list a + team model by its internal routing name while `/model/info` keeps its public name, so hide + both names to hide it on every route. + """ + return model_names + async def async_post_call_response_headers_hook( self, data: dict, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c0071aa7c81..a869150f7e8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -152,6 +152,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( @@ -11104,6 +11105,40 @@ class ProxyStartupEvent: #### API ENDPOINTS #### +async def _names_hidden_by_listing_callbacks( + user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] +) -> frozenset[str]: + hidden: Final = await proxy_logging_obj.hidden_by_listing_callbacks(user_api_key_dict, model_names) + if not hidden or llm_router is None: + return hidden + aliases: Final = llm_router.model_group_alias + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) + return hidden | frozenset( + alias + for alias in aliases + if (target := resolve_model_group_alias(aliases, alias)) is not None + and internal_to_public.get(target, target) in hidden + ) + + +async def _entries_kept_by_listing_callbacks( + entries: Sequence[tuple[str, str]], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, str], ...]: + hidden: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, tuple(response_id for response_id, _ in entries) + ) + if not hidden: + return tuple(entries) + return tuple(entry for entry in entries if entry[0] not in hidden) + + +async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_api_key_dict: UserAPIKeyAuth) -> bool: + listed_name: Final = _translate_model_name_for_response(deployment.model_dump(exclude_none=True)).get("model_name") + if not isinstance(listed_name, str): + return False + return listed_name in await _names_hidden_by_listing_callbacks(user_api_key_dict, (listed_name,)) + + @router.get("/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]) @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] @@ -11254,7 +11289,9 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - admin_entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + admin_entries: Final = await _entries_kept_by_listing_callbacks( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), user_api_key_dict + ) for response_id, lookup_id in admin_entries: model_info = create_model_info_response( model_id=lookup_id, @@ -11310,7 +11347,10 @@ async def model_list( # public name is what the client sees as the model id. model_data = [] entries: Final = alias_listing_entries( - TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases + await _entries_kept_by_listing_callbacks( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), user_api_key_dict + ), + caller_aliases, ) for response_id, lookup_id in entries: model_info = create_model_info_response( @@ -11404,13 +11444,24 @@ async def model_info( llm_router=llm_router, ) hidden_names: Final = blocked_names | unhealthy_names - if hidden_names: - all_models = [m for m in all_models if m not in hidden_names] + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, + tuple( + response_id + for response_id, _ in TeamModelNameTranslator.listing_entries( + tuple(m for m in all_models if m not in hidden_names), llm_router, settings + ) + ), + ) + if hidden_names or callback_hidden_names: + all_models = [ + m for m in all_models if m not in hidden_names and internal_to_public.get(m, m) not in callback_hidden_names + ] undiscoverable_names: Final = undiscoverable_model_names( all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id ) - internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) aliased_model_id: Final = alias_target( model_id, caller_alias_maps( @@ -15730,7 +15781,7 @@ async def model_info_v1( if litellm_model_id is not None: # user is trying to get specific model from litellm router deployment_info: Final = llm_router.get_deployment(model_id=litellm_model_id) - if deployment_info is None: + if deployment_info is None or await _deployment_hidden_by_listing_callbacks(deployment_info, user_api_key_dict): raise HTTPException( status_code=400, detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, @@ -15819,10 +15870,17 @@ async def model_info_v1( general_settings=general_settings, llm_router=llm_router, ) - visible_models: Final = discoverable_rows( + servable_rows: Final = discoverable_rows( (model for model in all_models if model.get("model_name") not in hidden_names), user_api_key_dict, ) + listed_names: Final = tuple( + dict.fromkeys(name for model in servable_rows if isinstance(name := model.get("model_name"), str)) + ) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks(user_api_key_dict, listed_names) + visible_models: Final = tuple( + model for model in servable_rows if model.get("model_name") not in callback_hidden_names + ) verbose_proxy_logger.debug("all_models: %s", visible_models) return _model_info_json_response(visible_models) @@ -15871,7 +15929,7 @@ async def model_deprecations( def _get_model_group_info( - llm_router: Router, all_models_str: list[str], model_group: str | None + llm_router: Router, all_models_str: Sequence[str], model_group: str | None ) -> list[ModelGroupInfoProxy]: model_groups: Final[list[ModelGroupInfoProxy]] = [] @@ -16104,23 +16162,34 @@ async def model_group_info( undiscoverable_group_names: Final = undiscoverable_model_names( all_models_str, llm_router, user_api_key_dict, user_api_key_dict.team_id ) - model_groups: list[ModelGroupInfoProxy] = _get_model_group_info( - llm_router=llm_router, - all_models_str=[name for name in all_models_str if name not in undiscoverable_group_names], - model_group=model_group, - ) + listed_group_names: Final = tuple(name for name in all_models_str if name not in undiscoverable_group_names) # Append A2A agents to model groups from litellm.proxy.agent_endpoints.model_list_helpers import ( append_agents_to_model_group, ) - model_groups = await append_agents_to_model_group( - model_groups=model_groups, + model_groups: Final = await append_agents_to_model_group( + model_groups=_get_model_group_info( + llm_router=llm_router, all_models_str=listed_group_names, model_group=model_group + ), user_api_key_dict=user_api_key_dict, ) + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) + public_group_names: Final = tuple( + internal_to_public.get(group.model_group, group.model_group) for group in model_groups + ) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, tuple(dict.fromkeys(public_group_names)) + ) - return {"data": model_groups} + return { + "data": [ + group + for group, public_name in zip(model_groups, public_group_names, strict=True) + if public_name not in callback_hidden_names + ] + } @router.get( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c617047fad9..b8cc30ad8a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -36,6 +36,7 @@ from typing import ( Final, Generic, Literal, + NoReturn, Optional, Protocol, TypeAlias, @@ -106,7 +107,7 @@ except ImportError: raise ImportError("backoff is not installed. Please install it via 'pip install backoff'") from fastapi import HTTPException, status -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm import litellm.litellm_core_utils @@ -1136,11 +1137,51 @@ class _CallbackCapabilities: # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) + listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) + + +def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: + leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__) + return any(hook_name in klass.__dict__ for klass in leaf_to_base) def _overrides_moderation_hook(callback: CustomLogger) -> bool: - leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__) - return any("async_moderation_hook" in klass.__dict__ for klass in leaf_to_base) + return _overrides_hook(callback, "async_moderation_hook") + + +_LISTED_MODEL_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class MalformedListingFilterReturn: + callback: str + tag: Literal["malformed_listing_filter_return"] = "malformed_listing_filter_return" + + +async def _names_kept_by_listing_callbacks( + callbacks: Sequence[CustomLogger], + user_api_key_dict: UserAPIKeyAuth, + model_names: tuple[str, ...], +) -> tuple[str, ...] | MalformedListingFilterReturn: + if not callbacks or not model_names: + return model_names + returned: Final = await callbacks[0].async_filter_listed_models(user_api_key_dict, model_names) + try: + kept: Final = frozenset(_LISTED_MODEL_NAMES.validate_python(returned)) + except ValidationError: + return MalformedListingFilterReturn(callback=type(callbacks[0]).__name__) + return await _names_kept_by_listing_callbacks( + callbacks[1:], user_api_key_dict, tuple(name for name in model_names if name in kept) + ) + + +def _raise_malformed_listing_filter_return(error: MalformedListingFilterReturn) -> NoReturn: + raise ProxyException( + message=f"{error.callback}.async_filter_listed_models must return a sequence of model names", + type=ProxyErrorTypes.internal_server_error, + param=None, + code=500, + ) class ProxyLogging: @@ -2808,6 +2849,9 @@ class ProxyLogging: has_moderation_override=has_moderation_override, iterator_overrides=tuple(iterator_overrides), resolved_callbacks=tuple(resolved_callbacks), + listed_models_filters=tuple( + callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") + ), ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. @@ -3715,6 +3759,18 @@ class ProxyLogging: verbose_proxy_logger.exception("Error in post_call_response_headers_hook: %s", str(e)) return merged_headers + async def hidden_by_listing_callbacks( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> frozenset[str]: + filters: Final = ProxyLogging._callback_capabilities().listed_models_filters + if not filters: + return frozenset() + candidates: Final = tuple(model_names) + kept: Final = await _names_kept_by_listing_callbacks(filters, user_api_key_dict, candidates) + if isinstance(kept, MalformedListingFilterReturn): + _raise_malformed_listing_filter_return(kept) + return frozenset(candidates).difference(kept) + @staticmethod def _build_litellm_call_info(data: dict, response: object) -> dict[str, object]: """ diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 636dc0f4d77..5175d92084c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -322,7 +322,9 @@ def test_get_proxy_model_info_shows_litellm_params_pricing_and_names_it_as_an_ov def test_get_proxy_model_info_names_config_model_info_pricing_as_an_override(monkeypatch, local_model_cost_map): """Pricing declared under ``model_info`` in config.yaml overrides the cost map too.""" info = _enriched_model_info( - monkeypatch, {"model": "openai/gpt-5.6"}, {"id": "dep-config", "db_model": False, "output_cost_per_token": 7e-06} + monkeypatch, + {"model": "openai/gpt-5.6"}, + {"id": "dep-config", "db_model": False, "output_cost_per_token": 7e-06}, ) assert info["pricing_overrides"] == ("output_cost_per_token",) assert info["output_cost_per_token"] == 7e-06 @@ -399,7 +401,9 @@ def test_model_info_reports_null_cost_for_unpriced_deployment_and_zero_for_decla def enriched_cost(model_name: str) -> tuple: deployment = router.get_model_list(model_name=model_name)[0] - info = proxy_server._enrich_model_info_with_litellm_data({**deployment, "model_info": dict(deployment["model_info"])})["model_info"] + info = proxy_server._enrich_model_info_with_litellm_data( + {**deployment, "model_info": dict(deployment["model_info"])} + )["model_info"] return info.get("input_cost_per_token"), info.get("output_cost_per_token") assert enriched_cost("vllm-unpriced") == (None, None) @@ -643,7 +647,6 @@ def model_group_info_router(monkeypatch): monkeypatch.setattr(proxy_server, "user_model", None) monkeypatch.setattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) monkeypatch.setattr(proxy_server, "user_api_key_cache", None) monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info) @@ -671,7 +674,9 @@ def test_model_group_info_proxy_admin_ignores_key_model_restriction( @pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) -def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role): +def test_model_group_info_proxy_admin_expands_wildcard_deployments( + client, auth_as, model_group_info_router, admin_role +): from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.model_checks import get_known_models_from_wildcard diff --git a/tests/test_litellm/proxy/test_model_list_callback_filter.py b/tests/test_litellm/proxy/test_model_list_callback_filter.py new file mode 100644 index 00000000000..00fbfee24ed --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_callback_filter.py @@ -0,0 +1,425 @@ +""" +Tests for `CustomLogger.async_filter_listed_models` on the model listing endpoints: +GET /v1/models (`model_list`, OpenAI and Anthropic shapes), GET /v1/models/{id} +(`model_info`), GET /v1/model/info (`model_info_v1`) and GET /model_group/info +(`model_group_info`). A registered callback that overrides the hook decides per +caller which of the names the route would list are kept; the rest disappear and +`/v1/models/{id}` answers 404 for them. +""" + +import json +from collections.abc import Sequence + +import pytest +from fastapi import HTTPException +from starlette.requests import Request + +import litellm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging + + +class _Gate(CustomLogger): + def __init__(self, hidden: frozenset[str] = frozenset(), extra: tuple[str, ...] = ()) -> None: + super().__init__() + self.hidden = hidden + self.extra = extra + self.seen: list[tuple[str, ...]] = [] + + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + self.seen.append(tuple(model_names)) + return [*(name for name in model_names if name not in self.hidden), *self.extra] + + +class _InferenceOnlyGate(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if data.get("model") == "restricted-model": + raise HTTPException(status_code=403, detail="not entitled to this model") + return data + + +class _RaisingGate(CustomLogger): + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + raise HTTPException(status_code=503, detail="entitlement service down") + + +class _ReversingGate(CustomLogger): + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + return list(reversed(model_names)) + + +class _StringReturningGate(CustomLogger): + async def async_filter_listed_models(self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str]) -> str: + return "open-model" + + +def _deployment(model_name: str, model: str = "openai/gpt-4o", **model_info): + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": f"{model_name}-id", **model_info}, + } + + +def _install_router(monkeypatch, *deployments, **router_kwargs) -> Router: + router = Router(model_list=list(deployments), **router_kwargs) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _register(monkeypatch, *callbacks: CustomLogger) -> None: + monkeypatch.setattr(litellm, "callbacks", list(callbacks)) + ProxyLogging._callback_capabilities_cache.clear() + + +@pytest.fixture +def two_model_router(monkeypatch) -> Router: + return _install_router(monkeypatch, _deployment("open-model"), _deployment("restricted-model")) + + +@pytest.fixture +def team_router(monkeypatch) -> Router: + return _install_router( + monkeypatch, + _deployment("gpt-4"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-gpt"), + _deployment("model_name_team1_def", team_id="team1", team_public_model_name="team-chat"), + ) + + +@pytest.fixture +def team_admin_privileges(monkeypatch) -> None: + from litellm.proxy.management_endpoints import common_utils + + async def _is_team_admin(**kwargs) -> bool: + return True + + monkeypatch.setattr(common_utils, "_user_has_admin_privileges", _is_team_admin) + + +def _non_admin(**kwargs) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.INTERNAL_USER, **kwargs) + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]) + + +def _team_member() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team1", + team_models=["model_name_team1_abc", "model_name_team1_def"], + models=["model_name_team1_abc", "model_name_team1_def"], + ) + + +def _anthropic_request() -> Request: + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/v1/models", + "query_string": b"", + "headers": [(b"anthropic-version", b"2023-06-01")], + } + ) + + +async def _v1_models(user_api_key_dict: UserAPIKeyAuth, **kwargs) -> list[str]: + response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, **kwargs) + return [m["id"] for m in response["data"]] + + +async def _v1_model_info_names(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + response = await proxy_server.model_info_v1(user_api_key_dict=user_api_key_dict) + return [row["model_name"] for row in json.loads(response.body)["data"]] + + +async def _model_groups(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + response = await proxy_server.model_group_info(user_api_key_dict=user_api_key_dict) + return [group.model_group for group in response["data"]] + + +async def _model_by_id_status(model_id: str, user_api_key_dict: UserAPIKeyAuth) -> int: + try: + response = await proxy_server.model_info(model_id=model_id, user_api_key_dict=user_api_key_dict) + except HTTPException as error: + return error.status_code + assert response["id"] == model_id + return 200 + + +@pytest.mark.asyncio +async def test_v1_models_lists_only_the_names_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_models(_non_admin()) == ["open-model"] + assert await _v1_models(_admin()) == ["open-model"] + assert await _v1_models(_non_admin(), request=_anthropic_request()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_models_scope_expand_applies_the_callback(two_model_router, team_admin_privileges, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_models(_non_admin(), scope="expand") == ["open-model"] + assert await _v1_models(_admin(), scope="expand") == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_models_by_id_answers_404_for_a_name_the_callback_leaves_out(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _model_by_id_status("restricted-model", _non_admin()) == 404 + assert await _model_by_id_status("open-model", _non_admin()) == 200 + + +@pytest.mark.asyncio +async def test_v1_model_info_lists_only_the_rows_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_model_info_names(_non_admin()) == ["open-model"] + assert await _v1_model_info_names(_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_model_group_info_lists_only_the_groups_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _model_groups(_non_admin()) == ["open-model"] + assert await _model_groups(_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_a_callback_without_the_hook_changes_no_listing(two_model_router, monkeypatch): + _register(monkeypatch, _InferenceOnlyGate()) + + assert await _v1_models(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_by_id_status("restricted-model", _non_admin()) == 200 + assert await _v1_model_info_names(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_groups(_non_admin()) == ["open-model", "restricted-model"] + + +@pytest.mark.asyncio +async def test_callback_cannot_add_a_name_it_was_not_offered(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(extra=("ghost-model",))) + + assert await _v1_models(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_by_id_status("ghost-model", _non_admin()) == 404 + + +@pytest.mark.asyncio +async def test_callbacks_narrow_in_registration_order(monkeypatch): + _install_router(monkeypatch, _deployment("a"), _deployment("b"), _deployment("c")) + first: _Gate = _Gate(hidden=frozenset({"a"})) + second: _Gate = _Gate(hidden=frozenset({"b"})) + _register(monkeypatch, first, second) + + assert await _v1_models(_non_admin()) == ["c"] + assert first.seen == [("a", "b", "c")] + assert second.seen == [("b", "c")] + + +@pytest.mark.asyncio +async def test_callback_sees_and_filters_team_models_by_their_public_name(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_models(_team_member()) == ["team-chat"] + assert await _model_by_id_status("team-gpt", _team_member()) == 404 + assert await _model_by_id_status("team-chat", _team_member()) == 200 + assert all("team-gpt" in seen and "model_name_team1_abc" not in seen for seen in gate.seen) + + +@pytest.mark.asyncio +async def test_callback_sees_public_team_names_on_every_listing_route(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_models(_team_member()) == ["team-chat"] + assert await _v1_model_info_names(_team_member()) == ["team-chat"] + assert await _model_groups(_team_member()) == ["model_name_team1_def"] + assert await _model_by_id_status("team-gpt", _team_member()) == 404 + assert len(gate.seen) == 4 + assert all(sorted(seen) == ["team-chat", "team-gpt"] for seen in gate.seen) + + +@pytest.mark.asyncio +async def test_router_alias_follows_its_hidden_target(monkeypatch): + _install_router( + monkeypatch, + _deployment("open-model"), + _deployment("restricted-model"), + model_group_alias={"mini": "restricted-model", "wide": "open-model"}, + ) + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert sorted(await _v1_models(_non_admin())) == ["open-model", "wide"] + assert sorted(await _v1_model_info_names(_non_admin())) == ["open-model", "wide"] + assert sorted(await _model_groups(_non_admin())) == ["open-model", "wide"] + + _register(monkeypatch, _Gate(hidden=frozenset({"mini"}))) + + assert sorted(await _v1_models(_non_admin())) == ["open-model", "restricted-model", "wide"] + + +@pytest.mark.asyncio +async def test_router_alias_of_a_team_model_follows_its_hidden_public_name(monkeypatch): + _install_router( + monkeypatch, + _deployment("gpt-4"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-gpt"), + model_group_alias={"team-alias": "model_name_team1_abc"}, + ) + caller: UserAPIKeyAuth = _non_admin( + user_id="u", + team_id="team1", + team_models=["model_name_team1_abc", "team-alias"], + models=["model_name_team1_abc", "team-alias"], + ) + _register(monkeypatch, _Gate(hidden=frozenset())) + assert sorted(await _v1_models(caller)) == ["team-alias", "team-gpt"] + + _register(monkeypatch, _Gate(hidden=frozenset({"team-gpt"}))) + assert await _v1_models(caller) == [] + assert await _model_groups(caller) == [] + + +@pytest.mark.asyncio +async def test_v1_model_info_offers_only_the_rows_the_caller_would_see(monkeypatch): + _install_router(monkeypatch, _deployment("open-model"), _deployment("hidden-model", discoverable=False)) + gate: _Gate = _Gate() + _register(monkeypatch, gate) + + assert await _v1_model_info_names(_non_admin()) == ["open-model"] + assert await _v1_model_info_names(_admin()) == ["open-model", "hidden-model"] + assert gate.seen == [("open-model",), ("open-model", "hidden-model")] + + +@pytest.mark.asyncio +async def test_listing_keeps_its_order_whatever_order_the_callback_returns(monkeypatch): + _install_router(monkeypatch, _deployment("a"), _deployment("b"), _deployment("c")) + _register(monkeypatch, _ReversingGate()) + + assert await _v1_models(_non_admin()) == ["a", "b", "c"] + assert await _v1_model_info_names(_non_admin()) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_a_callback_returning_a_string_is_an_error_not_an_empty_listing(two_model_router, monkeypatch): + _register(monkeypatch, _StringReturningGate()) + + with pytest.raises(ProxyException, match=r"_StringReturningGate\.async_filter_listed_models") as raised: + await _v1_models(_non_admin()) + assert raised.value.code == "500" + + +@pytest.mark.asyncio +async def test_alias_of_a_hidden_model_is_not_listed(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + caller = _non_admin(aliases={"mini": "restricted-model", "wide": "open-model"}) + + assert await _v1_models(caller) == ["open-model", "wide"] + assert await _model_by_id_status("mini", caller) == 404 + assert await _model_by_id_status("wide", caller) == 200 + + +@pytest.mark.asyncio +async def test_callback_error_reaches_the_caller(two_model_router, monkeypatch): + _register(monkeypatch, _RaisingGate()) + + with pytest.raises(HTTPException) as raised: + await _v1_models(_non_admin()) + assert raised.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_hidden_model_still_routes_for_direct_requests(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + assert "restricted-model" not in await _v1_models(_non_admin()) + + deployment = two_model_router.get_available_deployment( + model="restricted-model", messages=[{"role": "user", "content": "hi"}] + ) + assert deployment["model_name"] == "restricted-model" + + +@pytest.mark.asyncio +async def test_model_group_info_offers_a2a_agent_groups_to_the_callback(two_model_router, monkeypatch): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.types.agents import AgentResponse + + monkeypatch.setattr( + global_agent_registry, + "agent_list", + [AgentResponse(agent_id="agent-1", agent_name="helper", agent_card_params={})], + ) + caller = _non_admin(object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="p1", agents=["agent-1"])) + gate: _Gate = _Gate() + _register(monkeypatch, gate) + + assert await _model_groups(caller) == ["open-model", "restricted-model", "a2a/helper"] + assert gate.seen == [("open-model", "restricted-model", "a2a/helper")] + + _register(monkeypatch, _Gate(hidden=frozenset({"a2a/helper", "restricted-model"}))) + + assert await _model_groups(caller) == ["open-model"] + + +async def _v1_model_info_by_deployment_id(deployment_id: str, user_api_key_dict: UserAPIKeyAuth) -> int | list[str]: + try: + response = await proxy_server.model_info_v1(user_api_key_dict=user_api_key_dict, litellm_model_id=deployment_id) + except HTTPException as error: + return error.status_code + return [row["model_name"] for row in json.loads(response.body)["data"]] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_answers_like_an_unknown_id_for_a_hidden_model( + two_model_router, monkeypatch +): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_model_info_by_deployment_id("restricted-model-id", _non_admin()) == 400 + assert await _v1_model_info_by_deployment_id("no-such-id", _non_admin()) == 400 + assert await _v1_model_info_by_deployment_id("open-model-id", _non_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_offers_the_public_team_name(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_model_info_by_deployment_id("model_name_team1_abc-id", _team_member()) == 400 + assert await _v1_model_info_by_deployment_id("model_name_team1_def-id", _team_member()) == ["team-chat"] + assert gate.seen == [("team-gpt",), ("team-chat",)] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_offers_the_name_its_listing_shows_in_legacy_mode( + team_router, monkeypatch +): + monkeypatch.setattr(proxy_server, "general_settings", {"use_team_public_model_name": False}) + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_model_info_names(_team_member()) == ["team-chat"] + assert await _v1_model_info_by_deployment_id("model_name_team1_abc-id", _team_member()) == 400 + assert gate.seen[-1] == ("team-gpt",) From 248f0eb159c6c2788a28d4671104ffc4ce24904d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:59:11 +0000 Subject: [PATCH 027/154] ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests (#42903) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * build: point the local proxy unit targets at the nested tests/unit/proxy tree * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/classify_changes.sh | 2 +- .circleci/scripts/unit_selection.sh | 72 +++++++++ .circleci/tests.yml | 30 +++- .github/scripts/assert_ci_coverage.py | 1 - .github/workflows/test-unit-proxy-db.yml | 97 ++++------- .github/workflows/test-unit.yml | 8 +- Makefile | 10 +- litellm/llms/litellm_proxy/skills/README.md | 2 +- .../user_api_key_auth_code_coverage.py | 4 +- .../image_endpoints/test_azure_routes.py | 3 +- .../test_litellm/test_circleci_path_filter.py | 2 +- .../proxy/__init__.py} | 0 tests/unit/proxy/auth/__init__.py | 0 .../proxy/auth}/test_auth_checks.py | 0 .../test_default_end_user_budget_simple.py | 0 .../proxy/auth}/test_jwt.py | 0 .../auth}/test_models_fallback_endpoint.py | 0 .../auth}/test_multipart_bypass_repro.py | 0 .../proxy/auth}/test_proxy_routes.py | 0 .../proxy/auth}/test_user_api_key_auth.py | 0 tests/unit/proxy/common_utils/__init__.py | 0 .../common_utils}/test_check_batch_cost.py | 0 .../test_check_responses_cost.py | 0 .../test_proxy_encrypt_decrypt.py | 0 .../common_utils}/test_realtime_cache.py | 0 tests/unit/proxy/conftest.py | 150 ++++++++++++++++++ tests/unit/proxy/db/__init__.py | 0 .../proxy/db/db_transaction_queue/__init__.py | 0 .../test_e2e_pod_lock_manager.py | 0 .../proxy/db}/test_update_daily_tag_spend.py | 0 .../proxy/example_config_yaml/__init__.py | 0 .../example_config_yaml/aliases_config.yaml | 0 .../example_config_yaml/azure_config.yaml | 0 .../example_config_yaml/cache_no_params.yaml | 0 .../cache_with_params.yaml | 0 .../config_with_env_vars.yaml | 0 .../config_with_include.yaml | 0 .../config_with_missing_include.yaml | 0 .../config_with_multiple_includes.yaml | 0 .../example_config_yaml/included_models.yaml | 0 .../example_config_yaml/langfuse_config.yaml | 0 .../example_config_yaml/load_balancer.yaml | 0 .../example_config_yaml/models_file_1.yaml | 0 .../example_config_yaml/models_file_2.yaml | 0 .../opentelemetry_config.yaml | 0 .../example_config_yaml/simple_config.yaml | 0 tests/unit/proxy/google_endpoints/__init__.py | 0 .../test_gemini_agents_endpoints.py | 0 .../test_google_endpoint_routing.py | 0 .../test_google_gemini_proxy_request.py | 0 tests/unit/proxy/hooks/__init__.py | 0 .../proxy/hooks}/test_banned_keyword_list.py | 0 ...test_unit_test_max_model_budget_limiter.py | 0 .../proxy/management_endpoints/__init__.py | 0 .../test_jwt_key_mapping.py | 2 +- .../test_key_generate_prisma.py | 0 .../unit/proxy/management_helpers/__init__.py | 0 .../test_audit_logs_proxy.py | 0 tests/unit/proxy/middleware/__init__.py | 0 .../test_request_size_limit_middleware.py | 0 tests/unit/proxy/public_endpoints/__init__.py | 0 .../test_blog_posts_endpoint.py | 0 tests/unit/proxy/response_polling/__init__.py | 0 .../test_response_polling_handler.py | 0 tests/unit/proxy/spend_tracking/__init__.py | 0 .../test_search_api_logging.py | 0 .../proxy}/test_aproxy_startup.py | 0 tests/unit/proxy/test_configs/__init__.py | 0 .../proxy}/test_configs/custom_auth.py | 0 ...st_cloudflare_azure_with_cache_config.yaml | 0 .../proxy}/test_configs/test_config.yaml | 0 .../test_configs/test_config_custom_auth.yaml | 0 .../test_configs/test_config_no_auth.yaml | 0 .../test_configs/test_guardrails_config.yaml | 0 .../proxy}/test_custom_callback_input.py | 0 .../proxy}/test_custom_logger_s3_gcs.py | 0 .../proxy}/test_custom_tokenizer_bug.py | 0 .../proxy}/test_db_schema_changes.py | 0 .../test_deprecated_key_grace_period.py | 0 .../proxy}/test_get_favicon.py | 0 .../proxy}/test_get_image.py | 0 .../test_prisma_client_backoff_retry.py | 0 .../proxy}/test_prompt_test_endpoint.py | 0 .../proxy}/test_proxy_config_unit_test.py | 2 +- .../proxy}/test_proxy_custom_auth.py | 0 .../proxy}/test_proxy_reject_logging.py | 0 .../proxy}/test_proxy_server.py | 4 +- .../proxy}/test_proxy_setting_guardrails.py | 0 .../proxy}/test_proxy_token_counter.py | 0 .../proxy}/test_proxy_utils.py | 0 .../proxy}/test_reducto_ocr_route.py | 0 .../test_response_polling_pre_call_checks.py | 0 .../proxy}/test_server_root_path.py | 0 .../proxy}/test_ui_path_detection.py | 0 .../proxy}/test_unit_test_proxy_hooks.py | 0 .../proxy}/test_update_spend.py | 0 .../test_zero_cost_model_budget_bypass.py | 0 .../proxy}/vertex_key.json | 0 .../skills}/test_skills_db.py | 2 +- tests/unit/skills/test_skills_main.py | 2 +- 100 files changed, 305 insertions(+), 88 deletions(-) rename tests/{proxy_unit_tests/test_key_generate_dynamodb.py => unit/proxy/__init__.py} (100%) create mode 100644 tests/unit/proxy/auth/__init__.py rename tests/{proxy_unit_tests => unit/proxy/auth}/test_auth_checks.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_default_end_user_budget_simple.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_jwt.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_models_fallback_endpoint.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_multipart_bypass_repro.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_proxy_routes.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_user_api_key_auth.py (100%) create mode 100644 tests/unit/proxy/common_utils/__init__.py rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_check_batch_cost.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_check_responses_cost.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_proxy_encrypt_decrypt.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_realtime_cache.py (100%) create mode 100644 tests/unit/proxy/conftest.py create mode 100644 tests/unit/proxy/db/__init__.py create mode 100644 tests/unit/proxy/db/db_transaction_queue/__init__.py rename tests/{proxy_unit_tests => unit/proxy/db/db_transaction_queue}/test_e2e_pod_lock_manager.py (100%) rename tests/{proxy_unit_tests => unit/proxy/db}/test_update_daily_tag_spend.py (100%) create mode 100644 tests/unit/proxy/example_config_yaml/__init__.py rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/aliases_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/azure_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/cache_no_params.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/cache_with_params.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_env_vars.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_include.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_missing_include.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_multiple_includes.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/included_models.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/langfuse_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/load_balancer.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/models_file_1.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/models_file_2.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/opentelemetry_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/simple_config.yaml (100%) create mode 100644 tests/unit/proxy/google_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_gemini_agents_endpoints.py (100%) rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_google_endpoint_routing.py (100%) rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_google_gemini_proxy_request.py (100%) create mode 100644 tests/unit/proxy/hooks/__init__.py rename tests/{proxy_unit_tests => unit/proxy/hooks}/test_banned_keyword_list.py (100%) rename tests/{proxy_unit_tests => unit/proxy/hooks}/test_unit_test_max_model_budget_limiter.py (100%) create mode 100644 tests/unit/proxy/management_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/management_endpoints}/test_jwt_key_mapping.py (99%) rename tests/{proxy_unit_tests => unit/proxy/management_endpoints}/test_key_generate_prisma.py (100%) create mode 100644 tests/unit/proxy/management_helpers/__init__.py rename tests/{proxy_unit_tests => unit/proxy/management_helpers}/test_audit_logs_proxy.py (100%) create mode 100644 tests/unit/proxy/middleware/__init__.py rename tests/{proxy_unit_tests => unit/proxy/middleware}/test_request_size_limit_middleware.py (100%) create mode 100644 tests/unit/proxy/public_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/public_endpoints}/test_blog_posts_endpoint.py (100%) create mode 100644 tests/unit/proxy/response_polling/__init__.py rename tests/{proxy_unit_tests => unit/proxy/response_polling}/test_response_polling_handler.py (100%) create mode 100644 tests/unit/proxy/spend_tracking/__init__.py rename tests/{proxy_unit_tests => unit/proxy/spend_tracking}/test_search_api_logging.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_aproxy_startup.py (100%) create mode 100644 tests/unit/proxy/test_configs/__init__.py rename tests/{proxy_unit_tests => unit/proxy}/test_configs/custom_auth.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_cloudflare_azure_with_cache_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config_custom_auth.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config_no_auth.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_guardrails_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_callback_input.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_logger_s3_gcs.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_tokenizer_bug.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_db_schema_changes.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_deprecated_key_grace_period.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_get_favicon.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_get_image.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_prisma_client_backoff_retry.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_prompt_test_endpoint.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_config_unit_test.py (99%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_custom_auth.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_reject_logging.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_server.py (99%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_setting_guardrails.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_token_counter.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_utils.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_reducto_ocr_route.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_response_polling_pre_call_checks.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_server_root_path.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_ui_path_detection.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_unit_test_proxy_hooks.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_update_spend.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_zero_cost_model_budget_bypass.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/vertex_key.json (100%) rename tests/{proxy_unit_tests => unit/skills}/test_skills_db.py (98%) diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 01bc8290199..ad265a5e39f 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -31,7 +31,7 @@ while IFS= read -r file || [ -n "$file" ]; do case "$file" in model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json) has_cost_map=true ;; - tests/test_litellm/* | tests/proxy_unit_tests/*) : ;; + tests/test_litellm/* | tests/proxy_unit_tests/* | tests/unit/proxy/*) : ;; *) outside_cost_map_set=true ;; esac done diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 8d2b8a42691..2c60c5b1334 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,18 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + proxy-db-auth-checks + proxy-db-budgets + proxy-db-custom-logging + proxy-db-db-and-spend + proxy-db-endpoints-and-responses + proxy-db-guardrails-hooks + proxy-db-jwt-and-keys + proxy-db-key-generation + proxy-db-logging-misc + proxy-db-proxy-runtime + proxy-db-proxy-server-core + proxy-db-proxy-utils proxy-extras proxy-infra ) @@ -34,6 +46,66 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + proxy-db-auth-checks) + echo tests/unit/proxy/auth/test_auth_checks.py + echo tests/unit/proxy/auth/test_user_api_key_auth.py + echo tests/unit/proxy/test_deprecated_key_grace_period.py ;; + proxy-db-budgets) + echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py + echo tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py + echo tests/unit/proxy/test_zero_cost_model_budget_bypass.py ;; + proxy-db-custom-logging) + echo tests/unit/proxy/test_custom_callback_input.py + echo tests/unit/proxy/test_custom_logger_s3_gcs.py ;; + proxy-db-db-and-spend) + echo tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py + echo tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py + echo tests/unit/proxy/db/test_update_daily_tag_spend.py + echo tests/unit/proxy/test_db_schema_changes.py + echo tests/unit/proxy/test_prisma_client_backoff_retry.py + echo tests/unit/proxy/test_update_spend.py + echo tests/unit/skills/test_skills_db.py ;; + proxy-db-endpoints-and-responses) + echo tests/unit/proxy/auth/test_models_fallback_endpoint.py + echo tests/unit/proxy/common_utils/test_check_batch_cost.py + echo tests/unit/proxy/common_utils/test_check_responses_cost.py + echo tests/unit/proxy/common_utils/test_realtime_cache.py + echo tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py + echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py + echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py + echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py + echo tests/unit/proxy/response_polling/test_response_polling_handler.py + echo tests/unit/proxy/test_custom_tokenizer_bug.py + echo tests/unit/proxy/test_get_favicon.py + echo tests/unit/proxy/test_get_image.py + echo tests/unit/proxy/test_prompt_test_endpoint.py + echo tests/unit/proxy/test_reducto_ocr_route.py + echo tests/unit/proxy/test_response_polling_pre_call_checks.py + echo tests/unit/proxy/test_ui_path_detection.py ;; + proxy-db-guardrails-hooks) + echo tests/unit/proxy/hooks/test_banned_keyword_list.py + echo tests/unit/proxy/test_proxy_setting_guardrails.py + echo tests/unit/proxy/test_unit_test_proxy_hooks.py ;; + proxy-db-jwt-and-keys) + echo tests/unit/proxy/auth/test_jwt.py + echo tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py + echo tests/unit/proxy/test_proxy_custom_auth.py ;; + proxy-db-key-generation) echo tests/unit/proxy/management_endpoints/test_key_generate_prisma.py ;; + proxy-db-logging-misc) + echo tests/unit/proxy/management_helpers/test_audit_logs_proxy.py + echo tests/unit/proxy/spend_tracking/test_search_api_logging.py + echo tests/unit/proxy/test_proxy_reject_logging.py ;; + proxy-db-proxy-runtime) + echo tests/unit/proxy/auth/test_multipart_bypass_repro.py + echo tests/unit/proxy/auth/test_proxy_routes.py + echo tests/unit/proxy/middleware/test_request_size_limit_middleware.py + echo tests/unit/proxy/test_proxy_config_unit_test.py + echo tests/unit/proxy/test_proxy_token_counter.py + echo tests/unit/proxy/test_server_root_path.py ;; + proxy-db-proxy-server-core) + echo tests/unit/proxy/test_aproxy_startup.py + echo tests/unit/proxy/test_proxy_server.py ;; + proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; proxy-extras) echo tests/unit/litellm_proxy_extras ;; proxy-infra) echo tests/unit/gateway ;; *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 38fe44bb625..08e735637b7 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -317,7 +317,35 @@ workflows: reruns: 2 matrix: parameters: - flag: [enterprise-package, proxy-infra] + flag: + - enterprise-package + - proxy-infra + - proxy-db-auth-checks + - proxy-db-jwt-and-keys + - proxy-db-proxy-server-core + - proxy-db-proxy-runtime + - proxy-db-custom-logging + - proxy-db-logging-misc + - proxy-db-db-and-spend + - proxy-db-guardrails-hooks + - proxy-db-budgets + - proxy-db-endpoints-and-responses + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-proxy-db-proxy-utils + flag: proxy-db-proxy-utils + shards: 1 + reruns: 2 + dist: worksteal + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-proxy-db-key-generation + flag: proxy-db-key-generation + shards: 1 + workers: 0 + reruns: 2 base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - documentation diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index d8246225a3b..01a01b1034b 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -34,7 +34,6 @@ GLOB_CHARS = frozenset("*?") # tests has to be named by some shard or it runs nowhere. A child listed here is # itself decomposed one level deeper and is checked through its own entry. SHARDED_ROOTS: tuple[str, ...] = ( - "tests/proxy_unit_tests", "tests/test_litellm", "tests/test_litellm/proxy", ) diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 73015ac6e02..86b385d91a7 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -20,6 +20,12 @@ concurrency: # rather than alphabetical letter ranges. Adding a new test file means adding it # to whichever group it belongs to, not reshuffling slices. # +# `.circleci/tests.yml` runs each group's files on same-repo events under the +# `proxy-db-` Codecov flag; `.circleci/scripts/unit_selection.sh` holds +# the file lists. CircleCI does not build pull requests from forks, so `fork-flag` +# makes the shard run that list there. `test-path` keeps the files that still +# reach real providers and never left tests/proxy_unit_tests. +# # Design targets: # * Every shard runs in <= 7 minutes of wall-clock on the default runner. # Most of a shard's time is pytest plugin load + xdist worker imports + @@ -58,7 +64,7 @@ jobs: proxy-db: needs: assert-shard-coverage # Display only the semantic shard name in the checks UI instead of GHA's - # default "proxy-db (key-generation, tests/proxy_unit_tests/…, 0, loadscope, 20)" + # default "proxy-db (key-generation, tests/unit/proxy/…, 0, loadscope, 20)" # which includes every matrix field and gets truncated past the test-path. name: ${{ matrix.test-group }} permissions: @@ -71,132 +77,93 @@ jobs: include: # Must run serially — event-loop conflict with the logging worker. - test-group: key-generation - test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py" + test-path: "" + fork-flag: proxy-db-key-generation workers: 0 dist: loadscope timeout: 20 # ---- auth: split into 2 shards ---- - test-group: auth-checks - test-path: >- - tests/proxy_unit_tests/test_auth_checks.py - tests/proxy_unit_tests/test_user_api_key_auth.py - tests/proxy_unit_tests/test_deprecated_key_grace_period.py + test-path: "" + fork-flag: proxy-db-auth-checks workers: 4 dist: loadscope timeout: 15 - test-group: jwt-and-keys - test-path: >- - tests/proxy_unit_tests/test_jwt.py - tests/proxy_unit_tests/test_jwt_key_mapping.py - tests/proxy_unit_tests/test_proxy_custom_auth.py - tests/proxy_unit_tests/test_key_generate_dynamodb.py + test-path: "" + fork-flag: proxy-db-jwt-and-keys workers: 4 dist: loadscope timeout: 15 # ---- test_proxy_utils.py, single shard, worksteal distribution ---- - test-group: proxy-utils - test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + test-path: "" + fork-flag: proxy-db-proxy-utils workers: 4 dist: worksteal timeout: 15 # ---- proxy server: split into 2 shards ---- - test-group: proxy-server-core - test-path: >- - tests/proxy_unit_tests/test_proxy_server.py - tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py - tests/proxy_unit_tests/test_aproxy_startup.py + test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py" + fork-flag: proxy-db-proxy-server-core workers: 4 dist: loadscope timeout: 15 - test-group: proxy-runtime - test-path: >- - tests/proxy_unit_tests/test_proxy_config_unit_test.py - tests/proxy_unit_tests/test_proxy_routes.py - tests/proxy_unit_tests/test_server_root_path.py - tests/proxy_unit_tests/test_proxy_token_counter.py - tests/proxy_unit_tests/test_request_size_limit_middleware.py - tests/proxy_unit_tests/test_multipart_bypass_repro.py + test-path: "" + fork-flag: proxy-db-proxy-runtime workers: 4 dist: loadscope timeout: 15 # ---- logging: split into 2 shards ---- - test-group: custom-logging - test-path: >- - tests/proxy_unit_tests/test_custom_callback_input.py - tests/proxy_unit_tests/test_custom_logger_s3_gcs.py - tests/proxy_unit_tests/test_proxy_custom_logger.py + test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py" + fork-flag: proxy-db-custom-logging workers: 4 dist: loadscope timeout: 15 - test-group: logging-misc - test-path: >- - tests/proxy_unit_tests/test_proxy_reject_logging.py - tests/proxy_unit_tests/test_audit_logs_proxy.py - tests/proxy_unit_tests/test_search_api_logging.py + test-path: "" + fork-flag: proxy-db-logging-misc workers: 4 dist: loadscope timeout: 15 - test-group: db-and-spend - test-path: >- - tests/proxy_unit_tests/test_prisma_client_backoff_retry.py - tests/proxy_unit_tests/test_db_schema_changes.py - tests/proxy_unit_tests/test_e2e_pod_lock_manager.py - tests/proxy_unit_tests/test_skills_db.py - tests/proxy_unit_tests/test_update_daily_tag_spend.py - tests/proxy_unit_tests/test_update_spend.py - tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py + test-path: "" + fork-flag: proxy-db-db-and-spend workers: 4 dist: loadscope timeout: 15 # ---- guardrails + budget + hooks: split into 2 ---- - test-group: guardrails-hooks - test-path: >- - tests/proxy_unit_tests/test_proxy_setting_guardrails.py - tests/proxy_unit_tests/test_banned_keyword_list.py - tests/proxy_unit_tests/test_unit_test_proxy_hooks.py + test-path: "" + fork-flag: proxy-db-guardrails-hooks workers: 4 dist: loadscope timeout: 15 - test-group: budgets - test-path: >- - tests/proxy_unit_tests/test_default_end_user_budget_simple.py - tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py - tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py + test-path: "" + fork-flag: proxy-db-budgets workers: 4 dist: loadscope timeout: 15 - test-group: endpoints-and-responses - test-path: >- - tests/proxy_unit_tests/test_blog_posts_endpoint.py - tests/proxy_unit_tests/test_models_fallback_endpoint.py - tests/proxy_unit_tests/test_google_endpoint_routing.py - tests/proxy_unit_tests/test_google_gemini_proxy_request.py - tests/proxy_unit_tests/test_gemini_agents_endpoints.py - tests/proxy_unit_tests/test_get_favicon.py - tests/proxy_unit_tests/test_get_image.py - tests/proxy_unit_tests/test_reducto_ocr_route.py - tests/proxy_unit_tests/test_ui_path_detection.py - tests/proxy_unit_tests/test_prompt_test_endpoint.py - tests/proxy_unit_tests/test_check_batch_cost.py - tests/proxy_unit_tests/test_check_responses_cost.py - tests/proxy_unit_tests/test_response_polling_handler.py - tests/proxy_unit_tests/test_response_polling_pre_call_checks.py - tests/proxy_unit_tests/test_realtime_cache.py - tests/proxy_unit_tests/test_proxy_exception_mapping.py - tests/proxy_unit_tests/test_custom_tokenizer_bug.py + test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py" + fork-flag: proxy-db-endpoints-and-responses workers: 4 dist: loadscope timeout: 15 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} + fork-flag: ${{ matrix.fork-flag }} workers: ${{ matrix.workers }} reruns: 2 timeout-minutes: ${{ matrix.timeout }} diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 94de6040038..bf2e1602be8 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -31,10 +31,10 @@ concurrency: # number, so a partially-specified entry would fail the call rather than fall # back to the default. # -# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is -# already a matrix and carries a shard-coverage guard that reads that file by -# name. Folding it in here is a follow-up, together with generalising that guard -# into assert_ci_coverage.py. +# tests/unit/proxy keeps its own caller (test-unit-proxy-db.yml): it is already +# a matrix and carries a shard-coverage guard that reads that file by name. +# Folding it in here is a follow-up, together with generalising that guard into +# assert_ci_coverage.py. # # `fork-flag` names the `.circleci/tests.yml` job that now runs part of the # shard under the same Codecov flag. CircleCI does not build pull requests from diff --git a/Makefile b/Makefile index 6263b646c17..28daf589a23 100644 --- a/Makefile +++ b/Makefile @@ -51,8 +51,8 @@ help: @echo " make test-unit-core-utils - Run core utils tests (~32 files)" @echo " make test-unit-other - Run other tests (caching, responses, etc., ~69 files)" @echo " make test-unit-root - Run root-level tests (~34 files)" - @echo " make test-proxy-unit-a - Run proxy_unit_tests (a-o, ~20 files)" - @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" + @echo " make test-proxy-unit-a - Run tests/unit/proxy (a-o)" + @echo " make test-proxy-unit-b - Run tests/unit/proxy (p-z)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @@ -337,12 +337,12 @@ test-unit-other: install-test-deps test-unit-root: install-test-deps $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 -# Proxy unit tests (tests/proxy_unit_tests split alphabetically) +# Proxy unit tests (tests/unit/proxy split alphabetically) test-proxy-unit-a: install-test-deps - $(UV_RUN) pytest tests/proxy_unit_tests/test_[a-o]*.py --tb=short -vv -n 2 --durations=20 + $(UV_RUN) pytest tests/unit/proxy --ignore-glob='tests/unit/proxy/test_[p-z]*.py' --tb=short -vv -n 2 --durations=20 test-proxy-unit-b: install-test-deps - $(UV_RUN) pytest tests/proxy_unit_tests/test_[p-z]*.py --tb=short -vv -n 2 --durations=20 + $(UV_RUN) pytest tests/unit/proxy/test_[p-z]*.py tests/unit/skills --tb=short -vv -n 2 --durations=20 test-integration: install-test-deps $(UV_RUN) pytest tests/ -k "not test_litellm" diff --git a/litellm/llms/litellm_proxy/skills/README.md b/litellm/llms/litellm_proxy/skills/README.md index a896aa1166e..ccbd394cddc 100644 --- a/litellm/llms/litellm_proxy/skills/README.md +++ b/litellm/llms/litellm_proxy/skills/README.md @@ -369,7 +369,7 @@ model LiteLLM_SkillsTable { Run the tests: ```bash -pytest tests/proxy_unit_tests/test_skills_db.py -v +pytest tests/unit/skills/test_skills_db.py -v ``` Tests cover: diff --git a/tests/code_coverage_tests/user_api_key_auth_code_coverage.py b/tests/code_coverage_tests/user_api_key_auth_code_coverage.py index a9c2f8ef15f..2f221a7ebe7 100644 --- a/tests/code_coverage_tests/user_api_key_auth_code_coverage.py +++ b/tests/code_coverage_tests/user_api_key_auth_code_coverage.py @@ -31,11 +31,11 @@ def get_function_names_from_file(file_path): def get_all_functions_called_in_tests(base_dir): """ Returns a set of function names that are called in test functions - inside 'local_testing' and 'proxy_unit_tests' directories, + inside 'local_testing' and 'unit/proxy' directories, specifically in files containing the word 'router'. """ called_functions = set() - test_dirs = ["local_testing", "proxy_unit_tests"] + test_dirs = ["local_testing", "unit/proxy"] for test_dir in test_dirs: dir_path = os.path.join(base_dir, test_dir) diff --git a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py index 91fff717d25..46fe9a6f893 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py +++ b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py @@ -53,7 +53,8 @@ def client_no_auth(): config_fp = ( repo_root / "tests" - / "proxy_unit_tests" + / "unit" + / "proxy" / "test_configs" / "test_config_no_auth.yaml" ) diff --git a/tests/test_litellm/test_circleci_path_filter.py b/tests/test_litellm/test_circleci_path_filter.py index 84e2327057d..dcce7f57113 100644 --- a/tests/test_litellm/test_circleci_path_filter.py +++ b/tests/test_litellm/test_circleci_path_filter.py @@ -107,7 +107,7 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"] ), ( "cost-map-only", - ["model_prices_and_context_window.json", "tests/proxy_unit_tests/test_y.py"], + ["model_prices_and_context_window.json", "tests/unit/proxy/test_y.py"], "run", ), ( diff --git a/tests/proxy_unit_tests/test_key_generate_dynamodb.py b/tests/unit/proxy/__init__.py similarity index 100% rename from tests/proxy_unit_tests/test_key_generate_dynamodb.py rename to tests/unit/proxy/__init__.py diff --git a/tests/unit/proxy/auth/__init__.py b/tests/unit/proxy/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py similarity index 100% rename from tests/proxy_unit_tests/test_auth_checks.py rename to tests/unit/proxy/auth/test_auth_checks.py diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/unit/proxy/auth/test_default_end_user_budget_simple.py similarity index 100% rename from tests/proxy_unit_tests/test_default_end_user_budget_simple.py rename to tests/unit/proxy/auth/test_default_end_user_budget_simple.py diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/unit/proxy/auth/test_jwt.py similarity index 100% rename from tests/proxy_unit_tests/test_jwt.py rename to tests/unit/proxy/auth/test_jwt.py diff --git a/tests/proxy_unit_tests/test_models_fallback_endpoint.py b/tests/unit/proxy/auth/test_models_fallback_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_models_fallback_endpoint.py rename to tests/unit/proxy/auth/test_models_fallback_endpoint.py diff --git a/tests/proxy_unit_tests/test_multipart_bypass_repro.py b/tests/unit/proxy/auth/test_multipart_bypass_repro.py similarity index 100% rename from tests/proxy_unit_tests/test_multipart_bypass_repro.py rename to tests/unit/proxy/auth/test_multipart_bypass_repro.py diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/unit/proxy/auth/test_proxy_routes.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_routes.py rename to tests/unit/proxy/auth/test_proxy_routes.py diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_user_api_key_auth.py rename to tests/unit/proxy/auth/test_user_api_key_auth.py diff --git a/tests/unit/proxy/common_utils/__init__.py b/tests/unit/proxy/common_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py similarity index 100% rename from tests/proxy_unit_tests/test_check_batch_cost.py rename to tests/unit/proxy/common_utils/test_check_batch_cost.py diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py similarity index 100% rename from tests/proxy_unit_tests/test_check_responses_cost.py rename to tests/unit/proxy/common_utils/test_check_responses_cost.py diff --git a/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py b/tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py rename to tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py diff --git a/tests/proxy_unit_tests/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py similarity index 100% rename from tests/proxy_unit_tests/test_realtime_cache.py rename to tests/unit/proxy/common_utils/test_realtime_cache.py diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py new file mode 100644 index 00000000000..148751c33f2 --- /dev/null +++ b/tests/unit/proxy/conftest.py @@ -0,0 +1,150 @@ +# conftest.py + +import asyncio +import copy +import inspect +import warnings + +import pytest + + +import litellm +import litellm.proxy.proxy_server + + +# Top-level assignments of these types are the ones importlib.reload(litellm) +# would have effectively reset. We snapshot them at conftest import time and +# deep-copy the snapshot back before every test. +_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) + + +def _snapshot_mutable_state(module): + """Capture a per-module snapshot of primitive and collection attributes.""" + snapshot = {} + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception as exc: + warnings.warn( + f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", + stacklevel=2, + ) + continue + if value is None or isinstance(value, _SNAPSHOT_TYPES): + try: + snapshot[attr] = copy.deepcopy(value) + except Exception as exc: + warnings.warn( + f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + return snapshot + + +def _restore_mutable_state(module, snapshot): + for attr, default in snapshot.items(): + try: + setattr(module, attr, copy.deepcopy(default)) + except Exception as exc: + warnings.warn( + f"conftest: could not restore {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +def _collect_flushable_caches(): + """Return (module, attr) pairs whose values expose flush_cache().""" + targets = [] + for module in (litellm, litellm.proxy.proxy_server): + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + # Only instances — a class reference has an unbound flush_cache + # that can't be called without a self argument. + if inspect.isclass(value) or inspect.ismodule(value): + continue + if callable(getattr(value, "flush_cache", None)): + targets.append((module, attr)) + return targets + + +def _flush_caches(targets): + for module, attr in targets: + try: + value = getattr(module, attr) + except Exception: + continue + flush = getattr(value, "flush_cache", None) + if callable(flush): + try: + flush() + except Exception as exc: + warnings.warn( + f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +# Snapshot once at conftest import — these are the "clean" module states. +_LITELLM_STATE = _snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES = _collect_flushable_caches() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """Reset mutable module state on litellm and proxy_server before each test. + + Replaces a previous importlib.reload(litellm) approach that cost ~17s + per test (re-executing the full litellm __init__ import chain). + + What IS reset: + - Top-level module attributes of type list / dict / set / tuple + / str / int / float / bool / bytes, and None-valued attributes. + These cover callback lists, general_settings, master_key, + premium_user, prisma_client, etc. — anything the old reload() reset + by re-executing the module body. + - Any module-level object instance that exposes flush_cache() (the + DualCache and LLMClientCache family), which handles cache state + that can't round-trip through deepcopy because of internal locks. + + What is NOT reset: + - Class instances without flush_cache() (e.g. ProxyLogging, + JWTHandler, FastAPI routers, loggers). If a test mutates such an + instance in-place (setattr on the instance, appending to one of + its internal lists, etc.), the mutation will leak into later tests. + Use pytest's monkeypatch.setattr() or a local fixture for those + cases — don't rely on this autouse fixture to undo them. + """ + _restore_mutable_state(litellm, _LITELLM_STATE) + _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + _flush_caches(_FLUSHABLE_CACHES) + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/unit/proxy/db/__init__.py b/tests/unit/proxy/db/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/db/db_transaction_queue/__init__.py b/tests/unit/proxy/db/db_transaction_queue/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py b/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py similarity index 100% rename from tests/proxy_unit_tests/test_e2e_pod_lock_manager.py rename to tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/unit/proxy/db/test_update_daily_tag_spend.py similarity index 100% rename from tests/proxy_unit_tests/test_update_daily_tag_spend.py rename to tests/unit/proxy/db/test_update_daily_tag_spend.py diff --git a/tests/unit/proxy/example_config_yaml/__init__.py b/tests/unit/proxy/example_config_yaml/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/example_config_yaml/aliases_config.yaml b/tests/unit/proxy/example_config_yaml/aliases_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/aliases_config.yaml rename to tests/unit/proxy/example_config_yaml/aliases_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml b/tests/unit/proxy/example_config_yaml/azure_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/azure_config.yaml rename to tests/unit/proxy/example_config_yaml/azure_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/cache_no_params.yaml b/tests/unit/proxy/example_config_yaml/cache_no_params.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/cache_no_params.yaml rename to tests/unit/proxy/example_config_yaml/cache_no_params.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/cache_with_params.yaml b/tests/unit/proxy/example_config_yaml/cache_with_params.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/cache_with_params.yaml rename to tests/unit/proxy/example_config_yaml/cache_with_params.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_env_vars.yaml b/tests/unit/proxy/example_config_yaml/config_with_env_vars.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_env_vars.yaml rename to tests/unit/proxy/example_config_yaml/config_with_env_vars.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_include.yaml b/tests/unit/proxy/example_config_yaml/config_with_include.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_include.yaml rename to tests/unit/proxy/example_config_yaml/config_with_include.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_missing_include.yaml b/tests/unit/proxy/example_config_yaml/config_with_missing_include.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_missing_include.yaml rename to tests/unit/proxy/example_config_yaml/config_with_missing_include.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_multiple_includes.yaml b/tests/unit/proxy/example_config_yaml/config_with_multiple_includes.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_multiple_includes.yaml rename to tests/unit/proxy/example_config_yaml/config_with_multiple_includes.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/included_models.yaml b/tests/unit/proxy/example_config_yaml/included_models.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/included_models.yaml rename to tests/unit/proxy/example_config_yaml/included_models.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/langfuse_config.yaml b/tests/unit/proxy/example_config_yaml/langfuse_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/langfuse_config.yaml rename to tests/unit/proxy/example_config_yaml/langfuse_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/load_balancer.yaml b/tests/unit/proxy/example_config_yaml/load_balancer.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/load_balancer.yaml rename to tests/unit/proxy/example_config_yaml/load_balancer.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/models_file_1.yaml b/tests/unit/proxy/example_config_yaml/models_file_1.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/models_file_1.yaml rename to tests/unit/proxy/example_config_yaml/models_file_1.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/models_file_2.yaml b/tests/unit/proxy/example_config_yaml/models_file_2.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/models_file_2.yaml rename to tests/unit/proxy/example_config_yaml/models_file_2.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/opentelemetry_config.yaml b/tests/unit/proxy/example_config_yaml/opentelemetry_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/opentelemetry_config.yaml rename to tests/unit/proxy/example_config_yaml/opentelemetry_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/simple_config.yaml b/tests/unit/proxy/example_config_yaml/simple_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/simple_config.yaml rename to tests/unit/proxy/example_config_yaml/simple_config.yaml diff --git a/tests/unit/proxy/google_endpoints/__init__.py b/tests/unit/proxy/google_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_gemini_agents_endpoints.py b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py similarity index 100% rename from tests/proxy_unit_tests/test_gemini_agents_endpoints.py rename to tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py diff --git a/tests/proxy_unit_tests/test_google_endpoint_routing.py b/tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py similarity index 100% rename from tests/proxy_unit_tests/test_google_endpoint_routing.py rename to tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py similarity index 100% rename from tests/proxy_unit_tests/test_google_gemini_proxy_request.py rename to tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py diff --git a/tests/unit/proxy/hooks/__init__.py b/tests/unit/proxy/hooks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_banned_keyword_list.py b/tests/unit/proxy/hooks/test_banned_keyword_list.py similarity index 100% rename from tests/proxy_unit_tests/test_banned_keyword_list.py rename to tests/unit/proxy/hooks/test_banned_keyword_list.py diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py similarity index 100% rename from tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py rename to tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py diff --git a/tests/unit/proxy/management_endpoints/__init__.py b/tests/unit/proxy/management_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py similarity index 99% rename from tests/proxy_unit_tests/test_jwt_key_mapping.py rename to tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py index e95ed42013b..50b7a5c03fd 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py @@ -5,7 +5,7 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch # Add project root to sys.path -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))) from litellm.proxy.auth.user_api_key_auth import ( _resolve_jwt_to_virtual_key, diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py similarity index 100% rename from tests/proxy_unit_tests/test_key_generate_prisma.py rename to tests/unit/proxy/management_endpoints/test_key_generate_prisma.py diff --git a/tests/unit/proxy/management_helpers/__init__.py b/tests/unit/proxy/management_helpers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py similarity index 100% rename from tests/proxy_unit_tests/test_audit_logs_proxy.py rename to tests/unit/proxy/management_helpers/test_audit_logs_proxy.py diff --git a/tests/unit/proxy/middleware/__init__.py b/tests/unit/proxy/middleware/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_request_size_limit_middleware.py b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py similarity index 100% rename from tests/proxy_unit_tests/test_request_size_limit_middleware.py rename to tests/unit/proxy/middleware/test_request_size_limit_middleware.py diff --git a/tests/unit/proxy/public_endpoints/__init__.py b/tests/unit/proxy/public_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_blog_posts_endpoint.py b/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_blog_posts_endpoint.py rename to tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py diff --git a/tests/unit/proxy/response_polling/__init__.py b/tests/unit/proxy/response_polling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py similarity index 100% rename from tests/proxy_unit_tests/test_response_polling_handler.py rename to tests/unit/proxy/response_polling/test_response_polling_handler.py diff --git a/tests/unit/proxy/spend_tracking/__init__.py b/tests/unit/proxy/spend_tracking/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_search_api_logging.py b/tests/unit/proxy/spend_tracking/test_search_api_logging.py similarity index 100% rename from tests/proxy_unit_tests/test_search_api_logging.py rename to tests/unit/proxy/spend_tracking/test_search_api_logging.py diff --git a/tests/proxy_unit_tests/test_aproxy_startup.py b/tests/unit/proxy/test_aproxy_startup.py similarity index 100% rename from tests/proxy_unit_tests/test_aproxy_startup.py rename to tests/unit/proxy/test_aproxy_startup.py diff --git a/tests/unit/proxy/test_configs/__init__.py b/tests/unit/proxy/test_configs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_configs/custom_auth.py b/tests/unit/proxy/test_configs/custom_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_configs/custom_auth.py rename to tests/unit/proxy/test_configs/custom_auth.py diff --git a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/unit/proxy/test_configs/test_cloudflare_azure_with_cache_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml rename to tests/unit/proxy/test_configs/test_cloudflare_azure_with_cache_config.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config.yaml b/tests/unit/proxy/test_configs/test_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config.yaml rename to tests/unit/proxy/test_configs/test_config.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config_custom_auth.yaml b/tests/unit/proxy/test_configs/test_config_custom_auth.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config_custom_auth.yaml rename to tests/unit/proxy/test_configs/test_config_custom_auth.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml b/tests/unit/proxy/test_configs/test_config_no_auth.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml rename to tests/unit/proxy/test_configs/test_config_no_auth.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_guardrails_config.yaml b/tests/unit/proxy/test_configs/test_guardrails_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_guardrails_config.yaml rename to tests/unit/proxy/test_configs/test_guardrails_config.yaml diff --git a/tests/proxy_unit_tests/test_custom_callback_input.py b/tests/unit/proxy/test_custom_callback_input.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_callback_input.py rename to tests/unit/proxy/test_custom_callback_input.py diff --git a/tests/proxy_unit_tests/test_custom_logger_s3_gcs.py b/tests/unit/proxy/test_custom_logger_s3_gcs.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_logger_s3_gcs.py rename to tests/unit/proxy/test_custom_logger_s3_gcs.py diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_tokenizer_bug.py rename to tests/unit/proxy/test_custom_tokenizer_bug.py diff --git a/tests/proxy_unit_tests/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py similarity index 100% rename from tests/proxy_unit_tests/test_db_schema_changes.py rename to tests/unit/proxy/test_db_schema_changes.py diff --git a/tests/proxy_unit_tests/test_deprecated_key_grace_period.py b/tests/unit/proxy/test_deprecated_key_grace_period.py similarity index 100% rename from tests/proxy_unit_tests/test_deprecated_key_grace_period.py rename to tests/unit/proxy/test_deprecated_key_grace_period.py diff --git a/tests/proxy_unit_tests/test_get_favicon.py b/tests/unit/proxy/test_get_favicon.py similarity index 100% rename from tests/proxy_unit_tests/test_get_favicon.py rename to tests/unit/proxy/test_get_favicon.py diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/unit/proxy/test_get_image.py similarity index 100% rename from tests/proxy_unit_tests/test_get_image.py rename to tests/unit/proxy/test_get_image.py diff --git a/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py b/tests/unit/proxy/test_prisma_client_backoff_retry.py similarity index 100% rename from tests/proxy_unit_tests/test_prisma_client_backoff_retry.py rename to tests/unit/proxy/test_prisma_client_backoff_retry.py diff --git a/tests/proxy_unit_tests/test_prompt_test_endpoint.py b/tests/unit/proxy/test_prompt_test_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_prompt_test_endpoint.py rename to tests/unit/proxy/test_prompt_test_endpoint.py diff --git a/tests/proxy_unit_tests/test_proxy_config_unit_test.py b/tests/unit/proxy/test_proxy_config_unit_test.py similarity index 99% rename from tests/proxy_unit_tests/test_proxy_config_unit_test.py rename to tests/unit/proxy/test_proxy_config_unit_test.py index 5f236806685..2181c932586 100644 --- a/tests/proxy_unit_tests/test_proxy_config_unit_test.py +++ b/tests/unit/proxy/test_proxy_config_unit_test.py @@ -31,7 +31,7 @@ async def test_basic_reading_configs_from_files(): example_config_yaml_path = os.path.join(current_path, "example_config_yaml") # get all the files from example_config_yaml - files = os.listdir(example_config_yaml_path) + files = [f for f in os.listdir(example_config_yaml_path) if f.endswith((".yaml", ".yml"))] print(files) for file in files: diff --git a/tests/proxy_unit_tests/test_proxy_custom_auth.py b/tests/unit/proxy/test_proxy_custom_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_custom_auth.py rename to tests/unit/proxy/test_proxy_custom_auth.py diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/unit/proxy/test_proxy_reject_logging.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_reject_logging.py rename to tests/unit/proxy/test_proxy_reject_logging.py diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py similarity index 99% rename from tests/proxy_unit_tests/test_proxy_server.py rename to tests/unit/proxy/test_proxy_server.py index 5be27b3ad72..eae80f311d8 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -477,7 +477,7 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth): assert e.code == str(403) -from test_custom_callback_input import CompletionCustomHandler +from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler @mock_patch_acompletion() @@ -1114,7 +1114,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.team_endpoints import team_member_add -from test_key_generate_prisma import prisma_client +from tests.unit.proxy.management_endpoints.test_key_generate_prisma import prisma_client @pytest.fixture diff --git a/tests/proxy_unit_tests/test_proxy_setting_guardrails.py b/tests/unit/proxy/test_proxy_setting_guardrails.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_setting_guardrails.py rename to tests/unit/proxy/test_proxy_setting_guardrails.py diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_token_counter.py rename to tests/unit/proxy/test_proxy_token_counter.py diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_utils.py rename to tests/unit/proxy/test_proxy_utils.py diff --git a/tests/proxy_unit_tests/test_reducto_ocr_route.py b/tests/unit/proxy/test_reducto_ocr_route.py similarity index 100% rename from tests/proxy_unit_tests/test_reducto_ocr_route.py rename to tests/unit/proxy/test_reducto_ocr_route.py diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/unit/proxy/test_response_polling_pre_call_checks.py similarity index 100% rename from tests/proxy_unit_tests/test_response_polling_pre_call_checks.py rename to tests/unit/proxy/test_response_polling_pre_call_checks.py diff --git a/tests/proxy_unit_tests/test_server_root_path.py b/tests/unit/proxy/test_server_root_path.py similarity index 100% rename from tests/proxy_unit_tests/test_server_root_path.py rename to tests/unit/proxy/test_server_root_path.py diff --git a/tests/proxy_unit_tests/test_ui_path_detection.py b/tests/unit/proxy/test_ui_path_detection.py similarity index 100% rename from tests/proxy_unit_tests/test_ui_path_detection.py rename to tests/unit/proxy/test_ui_path_detection.py diff --git a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py b/tests/unit/proxy/test_unit_test_proxy_hooks.py similarity index 100% rename from tests/proxy_unit_tests/test_unit_test_proxy_hooks.py rename to tests/unit/proxy/test_unit_test_proxy_hooks.py diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/unit/proxy/test_update_spend.py similarity index 100% rename from tests/proxy_unit_tests/test_update_spend.py rename to tests/unit/proxy/test_update_spend.py diff --git a/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py similarity index 100% rename from tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py rename to tests/unit/proxy/test_zero_cost_model_budget_bypass.py diff --git a/tests/proxy_unit_tests/vertex_key.json b/tests/unit/proxy/vertex_key.json similarity index 100% rename from tests/proxy_unit_tests/vertex_key.json rename to tests/unit/proxy/vertex_key.json diff --git a/tests/proxy_unit_tests/test_skills_db.py b/tests/unit/skills/test_skills_db.py similarity index 98% rename from tests/proxy_unit_tests/test_skills_db.py rename to tests/unit/skills/test_skills_db.py index 8eb07a5ad48..20ffed6fec1 100644 --- a/tests/proxy_unit_tests/test_skills_db.py +++ b/tests/unit/skills/test_skills_db.py @@ -42,7 +42,7 @@ def create_skill_zip(skill_name: str): The zip file is automatically cleaned up after use. """ - test_dir = Path(__file__).parent.parent / "llm_translation" / "test_skills_data" + test_dir = Path(__file__).parents[2] / "llm_translation" / "test_skills_data" skill_dir = test_dir / skill_name # Create a zip file containing the skill directory diff --git a/tests/unit/skills/test_skills_main.py b/tests/unit/skills/test_skills_main.py index e1c66c8d9ea..71d65d45a08 100644 --- a/tests/unit/skills/test_skills_main.py +++ b/tests/unit/skills/test_skills_main.py @@ -30,7 +30,7 @@ def test_create_skill_forwards_description_and_instructions_from_top_level_kwarg def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: - """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + """The SDK convention (see tests/unit/skills/test_skills_db.py) nests them under extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" handler = MagicMock() monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) From ba776469916bcdb35ab238dec92f33ad428d724d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:07:48 +0000 Subject: [PATCH 028/154] ci: move provider-independent MCP tests into tests/unit and run mcp-integration from litellm-tests (#42904) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests * ci: move provider-independent MCP tests into tests/unit and run mcp-integration from litellm-tests * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * build: point the local proxy unit targets at the nested tests/unit/proxy tree * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/unit_selection.sh | 5 ++ .circleci/tests.yml | 21 +++++ .github/workflows/test-unit.yml | 1 + tests/unit/proxy/_experimental/__init__.py | 0 .../_experimental/mcp_server/__init__.py | 0 .../_experimental/mcp_server/conftest.py | 78 +++++++++++++++++++ .../test_mcp_auth_header_extraction.py | 0 .../mcp_server}/test_mcp_auth_priority.py | 0 .../mcp_server}/test_mcp_chat_completions.py | 0 .../mcp_server}/test_mcp_client_unit.py | 0 .../mcp_server}/test_mcp_logging.py | 0 .../mcp_server}/test_mcp_server.py | 0 .../mcp_server}/test_oauth2_mcp_config.yaml | 0 .../mcp_server}/test_openapi_spec_path_url.py | 0 .../mcp_server}/test_per_user_oauth_cache.py | 0 tests/unit/responses/__init__.py | 0 tests/unit/responses/mcp/__init__.py | 0 .../mcp}/test_aresponses_api_with_mcp.py | 0 18 files changed, 105 insertions(+) create mode 100644 tests/unit/proxy/_experimental/__init__.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/__init__.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/conftest.py rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_auth_header_extraction.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_auth_priority.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_chat_completions.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_client_unit.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_logging.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_server.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_oauth2_mcp_config.yaml (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_openapi_spec_path_url.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_per_user_oauth_cache.py (100%) create mode 100644 tests/unit/responses/__init__.py create mode 100644 tests/unit/responses/mcp/__init__.py rename tests/{mcp_tests => unit/responses/mcp}/test_aresponses_api_with_mcp.py (100%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 2c60c5b1334..f2ee7550df3 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,7 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + mcp-integration proxy-db-auth-checks proxy-db-budgets proxy-db-custom-logging @@ -46,6 +47,10 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + mcp-integration) + echo tests/unit/proxy/_experimental/mcp_server + echo tests/unit/responses/mcp + echo tests/mcp_tests/test_proxy_mcp_e2e.py ;; proxy-db-auth-checks) echo tests/unit/proxy/auth/test_auth_checks.py echo tests/unit/proxy/auth/test_user_api_key_auth.py diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 08e735637b7..264d7695a94 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -183,6 +183,9 @@ jobs: pull_request_url: type: string default: "" + legacy_mcp_peer: + type: boolean + default: false reruns: type: integer default: 0 @@ -200,6 +203,15 @@ jobs: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> - setup_test_deps + - when: + condition: << parameters.legacy_mcp_peer >> + steps: + - run: + name: Install MCP SDK1 peer + command: | + uv venv --python 3.12 .venv-mcp-peer + uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - run: name: "Run << parameters.flag >> shard" no_output_timeout: 20m @@ -215,6 +227,7 @@ jobs: rerun_args=(-p no:rerunfailures) if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") + if [ -n "${MCP_TEST_PEER_PYTHON:-}" ]; then test_env+=(MCP_TEST_PEER_PYTHON="$MCP_TEST_PEER_PYTHON"); fi set +e env -i "${test_env[@]}" \ uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml @@ -311,6 +324,14 @@ workflows: flag: [caching-local, proxy-extras, enterprise-routing] base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-mcp-integration + flag: mcp-integration + shards: 1 + workers: 2 + legacy_mcp_peer: true + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-<< matrix.flag >> shards: 1 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index bf2e1602be8..126a6e26e6f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -53,6 +53,7 @@ jobs: - shard: mcp-integration artifact-name: mcp-integration test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" + fork-flag: mcp-integration workers: 2 reruns: 0 timeout-minutes: 20 diff --git a/tests/unit/proxy/_experimental/__init__.py b/tests/unit/proxy/_experimental/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/_experimental/mcp_server/__init__.py b/tests/unit/proxy/_experimental/mcp_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/_experimental/mcp_server/conftest.py b/tests/unit/proxy/_experimental/mcp_server/conftest.py new file mode 100644 index 00000000000..d8b91e07467 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/conftest.py @@ -0,0 +1,78 @@ +import asyncio +import importlib + +import pytest + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + importlib.reload(litellm) + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + yield + + # Teardown code (executes after the yield point) + # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +@pytest.fixture(scope="function", autouse=True) +async def drain_logging_worker(): + """ + The logging queue is bound to the running loop, so anything left queued when a test's loop + goes away is carried onto the next test's loop and fires against its callbacks. + """ + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + yield + + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10) + except asyncio.TimeoutError: + pass + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests + + +@pytest.fixture +def config_only_mcp_manager_factory(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + class ConfigOnlyManager(MCPServerManager): + def initialize_tool_name_to_mcp_server_name_mapping(self): + return None + + return ConfigOnlyManager diff --git a/tests/mcp_tests/test_mcp_auth_header_extraction.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_header_extraction.py similarity index 100% rename from tests/mcp_tests/test_mcp_auth_header_extraction.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_header_extraction.py diff --git a/tests/mcp_tests/test_mcp_auth_priority.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_priority.py similarity index 100% rename from tests/mcp_tests/test_mcp_auth_priority.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_priority.py diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_chat_completions.py similarity index 100% rename from tests/mcp_tests/test_mcp_chat_completions.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_chat_completions.py diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py similarity index 100% rename from tests/mcp_tests/test_mcp_client_unit.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_logging.py similarity index 100% rename from tests/mcp_tests/test_mcp_logging.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_logging.py diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py similarity index 100% rename from tests/mcp_tests/test_mcp_server.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py diff --git a/tests/mcp_tests/test_oauth2_mcp_config.yaml b/tests/unit/proxy/_experimental/mcp_server/test_oauth2_mcp_config.yaml similarity index 100% rename from tests/mcp_tests/test_oauth2_mcp_config.yaml rename to tests/unit/proxy/_experimental/mcp_server/test_oauth2_mcp_config.yaml diff --git a/tests/mcp_tests/test_openapi_spec_path_url.py b/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py similarity index 100% rename from tests/mcp_tests/test_openapi_spec_path_url.py rename to tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/unit/proxy/_experimental/mcp_server/test_per_user_oauth_cache.py similarity index 100% rename from tests/mcp_tests/test_per_user_oauth_cache.py rename to tests/unit/proxy/_experimental/mcp_server/test_per_user_oauth_cache.py diff --git a/tests/unit/responses/__init__.py b/tests/unit/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/responses/mcp/__init__.py b/tests/unit/responses/mcp/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py similarity index 100% rename from tests/mcp_tests/test_aresponses_api_with_mcp.py rename to tests/unit/responses/mcp/test_aresponses_api_with_mcp.py From 3fa688223d97bb46345c8ed8032b583e03e29218 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:15:58 -0700 Subject: [PATCH 029/154] fix(vertex_ai): translate /v1/responses batch rows through the Responses-to-Chat bridge (#43042) * fix(vertex_ai): translate /v1/responses batch rows through the Responses-to-Chat bridge Vertex batch uploads treated every non-embeddings JSONL row as a chat completions body, so a /v1/responses row lost its input and reached GCS as a blank text part. Route detection now recognizes /v1/responses rows and bridges them to chat through the same Responses-to-Chat bridge the real-time path uses. That bridge call moves out of the Bedrock files transformation into a shared helper both providers call, forwarding the record's fields as sent, like real time, instead of validating them against the SDK TypedDicts whose required keys clients omit. * chore(batches): type the Vertex responses test helper and drop the quoted input cast * fix(batches): translate developer messages to system on Vertex and Bedrock batch rows like real time --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/base_llm/files/batch_records.py | 54 +++++++ litellm/llms/bedrock/files/transformation.py | 53 +----- .../llms/vertex_ai/files/transformation.py | 32 +++- .../test_bedrock_files_transformation.py | 43 +++++ .../test_vertex_ai_files_transformation.py | 153 ++++++++++++++++++ 5 files changed, 280 insertions(+), 55 deletions(-) create mode 100644 litellm/llms/base_llm/files/batch_records.py diff --git a/litellm/llms/base_llm/files/batch_records.py b/litellm/llms/base_llm/files/batch_records.py new file mode 100644 index 00000000000..6bb98456e69 --- /dev/null +++ b/litellm/llms/base_llm/files/batch_records.py @@ -0,0 +1,54 @@ +from collections.abc import Iterable, Mapping +from functools import cache +from types import MappingProxyType +from typing import Final, cast, get_type_hints + +from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIOptionalRequestParams + + +def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object]: + return MappingProxyType(dict(items)) + + +@cache +def _responses_request_keys() -> frozenset[str]: + return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams)) + + +def responses_batch_body_to_chat_body( + openai_request_body: Mapping[str, object], + custom_llm_provider: str | None = None, +) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict + """ + Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body. + + Batch providers translate chat bodies into their own request shape, so a Responses + record goes through the same Responses-to-Chat bridge the real-time path uses for + providers without a native Responses API: `input`, `instructions`, `max_output_tokens` + and the tool params translate identically in batch and real time. Like real time, the + record's fields are forwarded as sent instead of validated against the SDK TypedDicts, + whose required keys (a function tool's `strict`, an image part's `detail`) clients omit. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input: Final = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + model: Final = openai_request_body.get("model") + chat_input: Final = cast(str | ResponseInputParam, responses_input) # cast-ok: forwarded as sent + responses_request: Final = cast( # cast-ok: client-supplied fields forwarded verbatim, as real time does + ResponsesAPIOptionalRequestParams, + _frozen_mapping((key, value) for key, value in openai_request_body.items() if key in _responses_request_keys()), + ) + return LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # transformer declares a bare dict return + model=model if isinstance(model, str) else "", + input=chat_input, + responses_api_request=responses_request, + custom_llm_provider=custom_llm_provider, + metadata=openai_request_body.get("metadata"), + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 43faa7d79ea..a79f4de1e3d 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -8,7 +8,6 @@ from collections.abc import Iterable, Mapping, MutableMapping, Sequence from contextlib import suppress from dataclasses import dataclass from datetime import datetime -from functools import cache from itertools import chain from types import MappingProxyType from typing import Any, Final, Literal, TypeAlias, TypedDict @@ -17,7 +16,7 @@ from urllib.parse import quote, unquote, urlencode import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import ReadOnly from litellm._logging import verbose_logger @@ -41,7 +40,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, text_completion_prompt_to_messages, ) +from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -56,8 +57,6 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, - ResponseInputParam, - ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider @@ -130,22 +129,6 @@ class _S3UploadResponse(TypedDict, total=False): ContentLength: ReadOnly[int] -# JSONL batch records are untyped json, so the `/v1/responses` fields are -# validated into their concrete Responses API types before being handed to the -# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't -# define, which is what the bridge would ignore anyway. Built on first use -# rather than at import: `ResponseInputParam` is a deep union and only batch -# files carrying `/v1/responses` records need it. -@cache -def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: - return TypeAdapter(str | ResponseInputParam) - - -@cache -def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: - return TypeAdapter(ResponsesAPIOptionalRequestParams) - - class _BedrockS3RequestParams(AwsAuthParams): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -859,33 +842,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): Delegates to the same Responses-to-Chat bridge the real-time path uses for providers without a native Responses API (which is every Bedrock model), so `input`, `instructions`, `max_output_tokens` and the tool - params translate identically in batch and real time. The bridge always - emits a `tools` key; an empty one is dropped rather than shipped as an - empty array inside `modelInput`. + params translate identically in batch and real time. """ - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - responses_input: Final = openai_request_body.get("input") - if responses_input is None: - raise ValueError( - "Batch record for /v1/responses is missing required `input` field: " - f"model={openai_request_body.get('model', '')}" - ) - chat_body: Final[Mapping[str, object]] = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=openai_request_body.get("model", ""), - input=_responses_input_adapter().validate_python(responses_input), - responses_api_request=_responses_request_adapter().validate_python( - _frozen_mapping( - (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") - ) - ), - metadata=openai_request_body.get("metadata"), - ) - ) - return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + return responses_batch_body_to_chat_body(openai_request_body) @staticmethod def _transform_batch_body_to_chat_body( @@ -922,7 +881,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ from litellm.types.utils import LlmProviders - messages: Final = openai_request_body.get("messages", []) + messages: Final = map_developer_role_to_system_role(openai_request_body.get("messages", [])) optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 789b36ef3d0..dbb41b57348 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -35,7 +35,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, BaseFileUploadStream, @@ -529,21 +531,30 @@ def is_passthrough_batch_upload(create_file_data: Mapping[str, object], litellm_ return create_file_data.get("purpose") == "batch" and litellm_params.get("passthrough") is True -def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: +def _batch_entry_route_path(openai_entry: Mapping[str, object]) -> str: """ - Whether an OpenAI batch JSONL line targets the embeddings endpoint. + The route an OpenAI batch JSONL line targets, without query string or trailing slash. OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex has no equivalent per-line field, so the route decides which Vertex request shape the line has to be translated into. """ - url = openai_entry.get("url") + url: Final = openai_entry.get("url") if not isinstance(url, str): - return False - path = url.split("?")[0].rstrip("/") + return "" + return url.split("?")[0].rstrip("/") + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: + path: Final = _batch_entry_route_path(openai_entry) return path == "embeddings" or path.endswith("/embeddings") +def _is_responses_batch_entry(openai_entry: Mapping[str, object]) -> bool: + path: Final = _batch_entry_route_path(openai_entry) + return path == "responses" or path.endswith("/responses") + + def _openai_embedding_input_elements( embedding_input: GeminiEmbeddingInput, ) -> tuple[str | list[str], ...]: @@ -665,10 +676,15 @@ def _openai_batch_jsonl_entry_to_vertex_rows( return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) openai_request_body: Final = openai_entry.get("body") or {} + chat_request_body: Final = ( + responses_batch_body_to_chat_body(openai_request_body, custom_llm_provider="vertex_ai") + if _is_responses_batch_entry(openai_entry) + else openai_request_body + ) vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), + messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])), + model=chat_request_body.get("model", ""), + optional_params=map_openai_to_vertex_params(chat_request_body), custom_llm_provider="vertex_ai", litellm_params={}, cached_content=None, diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index d0921e68424..7a7159a1624 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1672,6 +1672,49 @@ class TestBedrockBatchNonChatEndpointRecords: assert "input" not in model_input assert "max_output_tokens" not in model_input + def test_anthropic_responses_record_accepts_a_function_tool_without_strict(self): + """Clients omit the SDK's required `strict`; the record is forwarded like real time, not validated.""" + parameters = {"type": "object", "properties": {"city": {"type": "string"}}} + model_input = self._transform( + { + "custom_id": "4a", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "Weather in Paris?", + "tools": [{"type": "function", "name": "get_weather", "parameters": parameters}], + }, + } + ) + + assert model_input["messages"][0]["content"] == [{"type": "text", "text": "Weather in Paris?"}] + tool = model_input["tools"][0] + function = tool.get("function", tool) + assert (function["name"], function.get("parameters", function.get("input_schema"))) == ("get_weather", parameters) + + @pytest.mark.parametrize( + ("url", "body"), + [ + ( + "/v1/responses", + {"input": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}]}, + ), + ( + "/v1/chat/completions", + {"messages": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}]}, + ), + ], + ids=["responses", "chat"], + ) + def test_anthropic_developer_role_becomes_the_system_prompt_like_real_time(self, url, body): + model_input = self._transform( + {"custom_id": "4c", "method": "POST", "url": url, "body": {"model": self.ANTHROPIC_MODEL, **body}} + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert [message["role"] for message in model_input["messages"]] == ["user"] + def test_responses_record_keeps_metadata(self): """`metadata` reaches the bridge, which reads it as its own kwarg.""" model_input = self._transform( diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 48464e79876..7434eae72a4 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -5,6 +5,7 @@ Includes tests for Vertex AI batch output transformation to OpenAI format. import json import urllib.parse +from collections.abc import Mapping from types import MappingProxyType from urllib.parse import parse_qs, urlparse @@ -1447,6 +1448,158 @@ class TestVertexEmbeddingsBatchInputTranslation: assert "content" in embeddings_row["request"] +def _responses_entry( + body: Mapping[str, object] | None = None, + custom_id: str = "resp-1", + url: str = "/v1/responses", +) -> dict[str, object]: + return { + "custom_id": custom_id, + "method": "POST", + "url": url, + "body": body + if body is not None + else {"model": "gemini-2.5-flash", "input": "What was the top headline in world news yesterday?"}, + } + + +class TestVertexResponsesBatchInputTranslation: + """ + /v1/responses batch lines carry `input`, not `messages`, so they go through the + Responses-to-Chat bridge before the Gemini translation instead of uploading as an + empty text part. + """ + + def test_string_input_becomes_the_user_prompt(self): + (row,) = _wrap_entries([_responses_entry()]) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "What was the top headline in world news yesterday?"}]} + ] + assert row["request"]["labels"]["litellm_custom_id"] == "resp-1" + + def test_instructions_and_input_items_map_like_real_time(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + } + ) + ] + ) + + request = row["request"] + assert request["system_instruction"] == {"parts": [{"text": "be terse"}]} + assert [content["role"] for content in request["contents"]] == ["user", "model", "user"] + assert request["contents"][-1]["parts"] == [{"text": "and 3+3?"}] + assert request["generationConfig"]["max_output_tokens"] == 32 + assert request["generationConfig"]["temperature"] == 0.2 + + def test_web_search_tool_keeps_the_prompt(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": "What was the top headline in world news yesterday?", + "tools": [{"type": "web_search"}], + } + ) + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "What was the top headline in world news yesterday?"}]} + ] + assert row["request"]["tools"] + + def test_sdk_optional_keys_are_not_required_like_real_time(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Weather in the pictured city?"}, + {"type": "input_image", "image_url": "https://example.com/paris.png"}, + ], + } + ], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + } + ) + ] + ) + + request = row["request"] + assert request["contents"][0]["parts"] == [ + {"text": "Weather in the pictured city?"}, + {"file_data": {"mime_type": "image/png", "file_uri": "https://example.com/paris.png"}}, + ] + assert request["tools"][0]["function_declarations"][0]["name"] == "get_weather" + + @pytest.mark.parametrize( + "url", + ["/v1/responses", "/v1/responses/", "/v1/responses?beta=1", "responses", "https://api.openai.com/v1/responses"], + ) + def test_route_spellings_are_all_responses(self, url): + (row,) = _wrap_entries([_responses_entry(url=url)]) + + assert row["request"]["contents"][0]["parts"] == [ + {"text": "What was the top headline in world news yesterday?"} + ] + + def test_missing_input_fails_the_upload(self): + with pytest.raises(ValueError, match="missing required `input` field"): + _wrap_entries([_responses_entry(body={"model": "gemini-2.5-flash"})]) + + @pytest.mark.parametrize( + "entry", + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}], + } + ), + { + "custom_id": "chat-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}], + }, + }, + ], + ids=["responses", "chat"], + ) + def test_developer_role_becomes_the_system_instruction_like_real_time(self, entry): + (row,) = _wrap_entries([entry]) + + request = row["request"] + assert request["system_instruction"] == {"parts": [{"text": "be terse"}]} + assert request["contents"] == [{"role": "user", "parts": [{"text": "ping"}]}] + + class TestVertexEmbeddingsBatchOutputTranslation: """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" From 1dc3b62dbc160c5f97bd8420829dcdd629e1e0c7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:26:44 -0700 Subject: [PATCH 030/154] fix(cost-map): add vertex priority prices for gemini-3-pro-image-preview and batch price for gemini-embedding-001 (#43069) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 13 +++++++++++++ model_prices_and_context_window.json | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9bec83f6b08..dc979cbd651 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26312,10 +26312,14 @@ "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26325,7 +26329,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -27848,6 +27854,7 @@ "gemini-embedding-001": { "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 1.2e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, "max_tokens": 2048, @@ -49492,10 +49499,14 @@ "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49505,7 +49516,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9bec83f6b08..dc979cbd651 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26312,10 +26312,14 @@ "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26325,7 +26329,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -27848,6 +27854,7 @@ "gemini-embedding-001": { "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 1.2e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, "max_tokens": 2048, @@ -49492,10 +49499,14 @@ "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49505,7 +49516,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" From 4584958574cd22b6f7f9b65aa6c92c58cdf940b6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:26:50 -0500 Subject: [PATCH 031/154] feat(agents): add optional per-agent kill switch webhook (#42841) --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/constants.py | 2 + litellm/proxy/_lazy_openapi_snapshot.json | 293 ++++++++++++++++++ litellm/proxy/_types.py | 4 +- .../proxy/agent_endpoints/agent_registry.py | 50 ++- litellm/proxy/agent_endpoints/endpoints.py | 90 +++++- litellm/proxy/agent_endpoints/kill_switch.py | 239 ++++++++++++++ litellm/proxy/schema.prisma | 1 + litellm/types/agents.py | 76 ++++- litellm/types/llms/custom_http.py | 1 + schema.prisma | 1 + .../agent_endpoints/test_agent_registry.py | 162 +++++++++- .../proxy/agent_endpoints/test_endpoints.py | 223 ++++++++++++- .../proxy/agent_endpoints/test_kill_switch.py | 248 +++++++++++++++ .../proxy/auth/test_route_checks.py | 1 + .../agents/_components/AgentFormKit.tsx | 5 + .../AgentKillSwitchDangerZone.test.tsx | 124 ++++++++ .../_components/AgentKillSwitchDangerZone.tsx | 147 +++++++++ .../_components/KillSwitchFormFields.tsx | 223 +++++++++++++ .../agents/_components/agent_config.ts | 23 ++ .../agents/_components/agent_form_fields.tsx | 9 + .../agents/_components/agent_info.test.tsx | 28 ++ .../agents/_components/agent_info.tsx | 11 +- .../_components/agent_type_utils.test.ts | 26 ++ .../agents/_components/agent_type_utils.ts | 2 + .../dynamic_agent_form_fields.test.ts | 65 ++++ .../_components/dynamic_agent_form_fields.tsx | 21 +- .../_components/kill_switch_config.test.ts | 118 +++++++ .../agents/_components/kill_switch_config.ts | 124 ++++++++ .../src/components/agents/types.ts | 3 + .../src/components/networking.tsx | 11 + .../AuditLogDrawer/AuditLogDrawer.tsx | 1 + .../components/view_logs/AuditLogsTable.tsx | 2 + .../view_logs/AuditLogsTableColumns.tsx | 9 +- .../src/contexts/PluginModeContext.tsx | 10 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 149 +++++++++ 37 files changed, 2485 insertions(+), 20 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql create mode 100644 litellm/proxy/agent_endpoints/kill_switch.py create mode 100644 tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/kill_switch_config.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/kill_switch_config.ts diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql new file mode 100644 index 00000000000..dd21ed644eb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "kill_switch" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/litellm/constants.py b/litellm/constants.py index 7b40f432446..807694c2f8a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -557,6 +557,8 @@ SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float( request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))) request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0 +AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000 # Patterns that indicate a localhost/internal URL in A2A agent cards that should be # replaced with the original base_url. This is a common misconfiguration where # developers deploy agents with development URLs in their agent cards. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0b43c3864ab..3d44315341b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2392,6 +2392,16 @@ ], "title": "Extra Headers" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "additionalProperties": true, "title": "Litellm Params", @@ -2561,6 +2571,221 @@ "title": "AgentKeySummary", "type": "object" }, + "AgentKillSwitchApiKeyAuth": { + "additionalProperties": false, + "properties": { + "api_key": { + "title": "Api Key", + "type": "string" + }, + "header_name": { + "default": "x-api-key", + "title": "Header Name", + "type": "string" + }, + "type": { + "const": "api_key", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "api_key" + ], + "title": "AgentKillSwitchApiKeyAuth", + "type": "object" + }, + "AgentKillSwitchBasicAuth": { + "additionalProperties": false, + "properties": { + "password": { + "title": "Password", + "type": "string" + }, + "type": { + "const": "basic", + "title": "Type", + "type": "string" + }, + "username": { + "title": "Username", + "type": "string" + } + }, + "required": [ + "type", + "username", + "password" + ], + "title": "AgentKillSwitchBasicAuth", + "type": "object" + }, + "AgentKillSwitchBearerAuth": { + "additionalProperties": false, + "properties": { + "token": { + "title": "Token", + "type": "string" + }, + "type": { + "const": "bearer", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "token" + ], + "title": "AgentKillSwitchBearerAuth", + "type": "object" + }, + "AgentKillSwitchConfig": { + "additionalProperties": false, + "description": "Webhook an admin fires to shut an agent down out of band. LiteLLM only\nmakes the call; whatever the endpoint does with it is the agent's business.", + "properties": { + "auth": { + "anyOf": [ + { + "discriminator": { + "mapping": { + "api_key": "#/components/schemas/AgentKillSwitchApiKeyAuth", + "basic": "#/components/schemas/AgentKillSwitchBasicAuth", + "bearer": "#/components/schemas/AgentKillSwitchBearerAuth" + }, + "propertyName": "type" + }, + "oneOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchBearerAuth" + }, + { + "$ref": "#/components/schemas/AgentKillSwitchApiKeyAuth" + }, + { + "$ref": "#/components/schemas/AgentKillSwitchBasicAuth" + } + ] + }, + { + "type": "null" + } + ], + "title": "Auth" + }, + "body": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Body" + }, + "headers": { + "additionalProperties": { + "type": "string" + }, + "title": "Headers", + "type": "object" + }, + "method": { + "default": "POST", + "enum": [ + "POST", + "PUT", + "PATCH", + "DELETE", + "GET" + ], + "title": "Method", + "type": "string" + }, + "query_params": { + "additionalProperties": { + "type": "string" + }, + "title": "Query Params", + "type": "object" + }, + "url": { + "title": "Url", + "type": "string" + } + }, + "required": [ + "url" + ], + "title": "AgentKillSwitchConfig", + "type": "object" + }, + "AgentKillSwitchResult": { + "properties": { + "agent_id": { + "title": "Agent Id", + "type": "string" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error" + }, + "method": { + "enum": [ + "POST", + "PUT", + "PATCH", + "DELETE", + "GET" + ], + "title": "Method", + "type": "string" + }, + "response_body": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Response Body" + }, + "status_code": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Status Code" + }, + "url": { + "title": "Url", + "type": "string" + } + }, + "required": [ + "agent_id", + "url", + "method" + ], + "title": "AgentKillSwitchResult", + "type": "object" + }, "AgentMakePublicResponse": { "properties": { "message": { @@ -2775,6 +3000,16 @@ ], "title": "Keys" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "anyOf": [ { @@ -3569,6 +3804,16 @@ ], "title": "Extra Headers" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "additionalProperties": true, "title": "Litellm Params", @@ -4331,6 +4576,54 @@ ] } }, + "/v1/agents/{agent_id}/kill_switch": { + "post": { + "description": "Fire the agent's configured kill switch webhook. Proxy admin only.\n\nLiteLLM only makes the configured HTTP call and reports what came back; it\ndoes not change the agent's state in LiteLLM. Returns 200 when the webhook\nanswered 2xx, 502 with the same result body otherwise. Every attempt is\nwritten to the audit log as a `kill_switch_fired` row against the agent.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch\" \\\n -H \"Authorization: Bearer \"\n```", + "operationId": "trigger_agent_kill_switch_v1_agents__agent_id__kill_switch_post", + "parameters": [ + { + "in": "path", + "name": "agent_id", + "required": true, + "schema": { + "title": "Agent Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AgentKillSwitchResult" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Trigger Agent Kill Switch", + "tags": [ + "agents" + ] + } + }, "/v1/agents/{agent_id}/make_public": { "post": { "description": "Make an agent publicly discoverable\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/make_public\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\"\n```\n\nExample Response:\n```json\n{\n \"agent_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"agent_name\": \"my-custom-agent\",\n \"litellm_params\": {\n \"make_public\": true\n },\n \"agent_card_params\": {...},\n \"created_at\": \"2025-11-15T10:30:00Z\",\n \"updated_at\": \"2025-11-15T10:35:00Z\",\n \"created_by\": \"user123\",\n \"updated_by\": \"user123\"\n}\n```", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6de36f8423..12b4d4b2412 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -238,6 +238,7 @@ class LitellmTableNames(str, enum.Enum): CONFIG_TABLE_NAME = "LiteLLM_Config" SSO_CONFIG_TABLE_NAME = "LiteLLM_SSOConfig" UI_SETTINGS_TABLE_NAME = "LiteLLM_UISettings" + AGENT_TABLE_NAME = "LiteLLM_AgentsTable" class Litellm_EntityType(enum.Enum): @@ -578,6 +579,7 @@ class LiteLLMRoutes(enum.Enum): "/v1/agents/{agent_id}", "/v1/agents/make_public", "/v1/agents/{agent_id}/make_public", + "/v1/agents/{agent_id}/kill_switch", ) # Backwards-compat union — virtual keys may be configured with @@ -3688,7 +3690,7 @@ from litellm.models.spend_logs import ( # noqa: E402 ) from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402 -AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated"] +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated", "kill_switch_fired"] class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 3d56c2b5326..3e775d7648e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -13,13 +13,14 @@ import litellm from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.proxy.agent_endpoints.kill_switch import restore_kill_switch from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository -from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest +from litellm.types.agents import AgentConfig, AgentKillSwitchConfig, AgentResponse, PatchAgentRequest if TYPE_CHECKING: from prisma import models as prisma_models @@ -31,6 +32,10 @@ class AgentObjectPermissionRecord(Protocol): def dict(self) -> dict[str, object]: ... +class AgentIdWhere(TypedDict): + agent_id: ReadOnly[str] + + class AgentRecordDump(TypedDict): agent_id: str agent_name: str @@ -38,6 +43,7 @@ class AgentRecordDump(TypedDict): agent_card_params: dict[str, object] static_headers: dict[str, str] | None extra_headers: list[str] | None + kill_switch: ReadOnly[AgentKillSwitchConfig | None] access_group_ids: ReadOnly[Sequence[str] | None] object_permission: dict[str, object] | None spend: float @@ -70,6 +76,9 @@ class AgentRecord(Protocol): @property def access_group_ids(self) -> Sequence[str] | None: ... + @property + def kill_switch(self) -> Mapping[str, object] | None: ... + @property def spend(self) -> float: ... @@ -211,6 +220,29 @@ def parse_agent_litellm_params(value: object) -> Mapping[str, object]: return _EMPTY_LITELLM_PARAMS +_KILL_SWITCH_ADAPTER: Final[TypeAdapter[AgentKillSwitchConfig | None]] = TypeAdapter(AgentKillSwitchConfig | None) + + +def parse_agent_kill_switch(value: object) -> AgentKillSwitchConfig | None: + if value is None: + return None + try: + if isinstance(value, str): + return _KILL_SWITCH_ADAPTER.validate_json(value) + return _KILL_SWITCH_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def serialize_agent_kill_switch(incoming: object, existing: object) -> str: + """prisma-client-py drops ``None`` from update data, so a cleared kill switch is stored as the JSON literal + ``null`` (read back as ``None``), the same convention ``memory_endpoints`` uses for ``Json?`` columns.""" + restored: Final = restore_kill_switch( + _KILL_SWITCH_ADAPTER.validate_python(incoming), parse_agent_kill_switch(existing) + ) + return safe_dumps(restored.model_dump() if restored is not None else None) + + _MISSING_AGENT_PARAM: Final = object() _RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10 @@ -293,6 +325,12 @@ def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) +def _patched_kill_switch(agent: PatchAgentRequest, existing: object) -> Mapping[str, object]: + if "kill_switch" not in agent: + return MappingProxyType({}) + return MappingProxyType({"kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), existing)}) + + def _restore_redacted_litellm_params( incoming: Mapping[str, object], existing: Mapping[str, object], @@ -531,6 +569,7 @@ class AgentRegistry: "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, + "kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), None), "created_by": created_by, "updated_by": created_by, "created_at": datetime.now(timezone.utc), @@ -613,7 +652,10 @@ class AgentRegistry: existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)} + update_data: Final[dict[str, object]] = { + **_patched_access_group_ids(agent), + **_patched_kill_switch(agent, existing_agent.get("kill_switch")), + } if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if "litellm_params" in agent: @@ -716,6 +758,9 @@ class AgentRegistry: ) extra_headers_val_u: Final = agent.get("extra_headers") or [] access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) + kill_switch_val_u: Final = serialize_agent_kill_switch( + agent.get("kill_switch"), existing_row.kill_switch if existing_row is not None else None + ) update_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -723,6 +768,7 @@ class AgentRegistry: "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, + "kill_switch": kill_switch_val_u, "access_group_ids": access_group_ids_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index aa8979a73c6..28c82a715e0 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -33,6 +33,8 @@ from litellm.proxy.a2a.agent_card import ( normalize_protocol_version, ) from litellm.proxy.agent_endpoints.agent_registry import ( + AgentIdWhere, + parse_agent_kill_switch, parse_agent_litellm_params, redact_sensitive_agent_litellm_params, ) @@ -45,6 +47,15 @@ from litellm.proxy.agent_endpoints.agent_search import ( search_agents, ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents +from litellm.proxy.agent_endpoints.kill_switch import ( + KillSwitchAuditLogWriter, + KillSwitchHttpClient, + build_kill_switch_audit_log, + default_kill_switch_audit_log_writer, + default_kill_switch_http_client, + fire_kill_switch, + redact_kill_switch, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -53,6 +64,8 @@ from litellm.types.agents import ( AgentCard, AgentConfig, AgentKeySummary, + AgentKillSwitchConfig, + AgentKillSwitchResult, AgentMakePublicResponse, AgentResponse, MakeAgentsPublicRequest, @@ -160,9 +173,10 @@ def _redact_sensitive_agent_fields( ) -> list[AgentResponse]: """ Return copies of the given agents with credential-bearing litellm_params - values replaced by a fixed marker (never returned to ANY caller, - admin included) and, for non-admin callers, virtual-key and header - fields stripped entirely. The original objects are not modified. + values and kill-switch auth secrets replaced by a fixed marker (never + returned to ANY caller, admin included) and, for non-admin callers, + virtual-key, header and kill-switch fields stripped entirely. The original + objects are not modified. """ redacted: Final[list[AgentResponse]] = [] for agent in agents: @@ -171,8 +185,10 @@ def _redact_sensitive_agent_fields( copy.static_headers = None copy.extra_headers = None copy.keys = None + copy.kill_switch = None if copy.litellm_params: copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params) + copy.kill_switch = redact_kill_switch(copy.kill_switch) redacted.append(copy) return redacted @@ -872,6 +888,74 @@ async def delete_agent( raise HTTPException(status_code=500, detail=str(e)) +@router.post( + "/v1/agents/{agent_id}/kill_switch", + tags=["[beta] A2A Agents"], # mutable-ok: fastapi types tags as list[str | Enum] + dependencies=(Depends(user_api_key_auth),), + response_model=AgentKillSwitchResult, +) +async def trigger_agent_kill_switch( + agent_id: str, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + http_client: Annotated[KillSwitchHttpClient, Depends(default_kill_switch_http_client)], + audit_log_writer: Annotated[KillSwitchAuditLogWriter, Depends(default_kill_switch_audit_log_writer)], +): + """ + Fire the agent's configured kill switch webhook. Proxy admin only. + + LiteLLM only makes the configured HTTP call and reports what came back; it + does not change the agent's state in LiteLLM. Returns 200 when the webhook + answered 2xx, 502 with the same result body otherwise. Every attempt is + written to the audit log as a `kill_switch_fired` row against the agent. + + Example Request: + ```bash + curl -X POST "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch" \\ + -H "Authorization: Bearer " + ``` + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + await check_feature_access_for_user(user_api_key_dict, "agents") + _check_agent_management_permission(user_api_key_dict) + + resolved: Final = await _resolve_agent_kill_switch(agent_id) + if resolved is None: + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") + resolved_agent_id, config = resolved + if config is None: + raise HTTPException(status_code=400, detail=f"Agent with ID {agent_id} has no kill_switch configured") + + result: Final = await fire_kill_switch(agent_id=resolved_agent_id, config=config, http_client=http_client) + await audit_log_writer( + build_kill_switch_audit_log( + result=result, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + ) + if not result.succeeded: + raise HTTPException(status_code=502, detail=result.model_dump()) + return result + + +async def _resolve_agent_kill_switch(agent_id: str) -> tuple[str, AgentKillSwitchConfig | None] | None: + """The DB row wins over this replica's in-memory registry so a trigger never fires a webhook another + replica has since changed; config.yaml agents have no row and fall back to the registry.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + where: Final[AgentIdWhere] = {"agent_id": agent_id} + row: Final = await agents_table(prisma_client).find_unique(where=where) + if row is not None: + return row.agent_id, parse_agent_kill_switch(row.kill_switch) + + agent: Final = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) + if agent is None: + return None + return agent.agent_id, agent.kill_switch + + @router.post( "/v1/agents/{agent_id}/make_public", tags=["[beta] A2A Agents"], diff --git a/litellm/proxy/agent_endpoints/kill_switch.py b/litellm/proxy/agent_endpoints/kill_switch.py new file mode 100644 index 00000000000..8b3f64e74ee --- /dev/null +++ b/litellm/proxy/agent_endpoints/kill_switch.py @@ -0,0 +1,239 @@ +from base64 import b64encode +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias + +import httpx +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.constants import ( + AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS, + AGENT_KILL_SWITCH_TIMEOUT_SECONDS, + REDACTED_BY_LITELM_STRING, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params arg is a bare dict in http_handler +) +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update, get_audit_log_changed_by +from litellm.types.agents import ( + AgentKillSwitchApiKeyAuth, + AgentKillSwitchAuth, + AgentKillSwitchBasicAuth, + AgentKillSwitchBearerAuth, + AgentKillSwitchConfig, + AgentKillSwitchResult, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + + +def _with_auth(config: AgentKillSwitchConfig, auth: AgentKillSwitchAuth) -> AgentKillSwitchConfig: + return AgentKillSwitchConfig( + url=config.url, + method=config.method, + headers=config.headers, + query_params=config.query_params, + body=config.body, + auth=auth, + ) + + +def redact_kill_switch(config: AgentKillSwitchConfig | None) -> AgentKillSwitchConfig | None: + if config is None or config.auth is None: + return config + return _with_auth(config, _redact_auth(config.auth)) + + +def _redact_auth(auth: AgentKillSwitchAuth) -> AgentKillSwitchAuth: + match auth: + case AgentKillSwitchBearerAuth(): + return AgentKillSwitchBearerAuth(type="bearer", token=REDACTED_BY_LITELM_STRING) + case AgentKillSwitchApiKeyAuth(): + return AgentKillSwitchApiKeyAuth( + type="api_key", header_name=auth.header_name, api_key=REDACTED_BY_LITELM_STRING + ) + case AgentKillSwitchBasicAuth(): + return AgentKillSwitchBasicAuth(type="basic", username=auth.username, password=REDACTED_BY_LITELM_STRING) + case _: + assert_never(auth) + + +def restore_kill_switch( + incoming: AgentKillSwitchConfig | None, + existing: AgentKillSwitchConfig | None, +) -> AgentKillSwitchConfig | None: + """Put the stored secret back behind an auth field echoed as the redaction + marker; a marker with no stored secret of the same auth type becomes "".""" + if incoming is None or incoming.auth is None: + return incoming + existing_auth: Final = existing.auth if existing is not None else None + return _with_auth(incoming, _restore_auth(incoming.auth, existing_auth)) + + +def _restore_secret(incoming_value: str, existing_value: str | None) -> str: + if incoming_value != REDACTED_BY_LITELM_STRING: + return incoming_value + return existing_value if existing_value is not None else "" + + +def _restore_auth(incoming: AgentKillSwitchAuth, existing: AgentKillSwitchAuth | None) -> AgentKillSwitchAuth: + match incoming: + case AgentKillSwitchBearerAuth(): + stored_token: Final = existing.token if isinstance(existing, AgentKillSwitchBearerAuth) else None + return AgentKillSwitchBearerAuth(type="bearer", token=_restore_secret(incoming.token, stored_token)) + case AgentKillSwitchApiKeyAuth(): + stored_key: Final = existing.api_key if isinstance(existing, AgentKillSwitchApiKeyAuth) else None + return AgentKillSwitchApiKeyAuth( + type="api_key", + header_name=incoming.header_name, + api_key=_restore_secret(incoming.api_key, stored_key), + ) + case AgentKillSwitchBasicAuth(): + stored_password: Final = existing.password if isinstance(existing, AgentKillSwitchBasicAuth) else None + return AgentKillSwitchBasicAuth( + type="basic", + username=incoming.username, + password=_restore_secret(incoming.password, stored_password), + ) + case _: + assert_never(incoming) + + +@dataclass(frozen=True, slots=True) +class KillSwitchRequest: + method: str + url: str + headers: Mapping[str, str] + json_body: Mapping[str, object] | None + + +def _auth_headers(auth: AgentKillSwitchAuth | None) -> Mapping[str, str]: + match auth: + case None: + return MappingProxyType({}) + case AgentKillSwitchBearerAuth(): + return MappingProxyType({"Authorization": f"Bearer {auth.token}"}) + case AgentKillSwitchApiKeyAuth(): + return MappingProxyType({auth.header_name: auth.api_key}) + case AgentKillSwitchBasicAuth(): + credentials: Final = b64encode(f"{auth.username}:{auth.password}".encode()).decode() + return MappingProxyType({"Authorization": f"Basic {credentials}"}) + case _: + assert_never(auth) + + +def build_kill_switch_request(config: AgentKillSwitchConfig) -> KillSwitchRequest: + url: Final = httpx.URL(config.url).copy_merge_params(config.query_params) + return KillSwitchRequest( + method=config.method, + url=str(url), + headers=MappingProxyType({**config.headers, **_auth_headers(config.auth)}), + json_body=config.body, + ) + + +class KillSwitchHttpClient(Protocol): + def build_request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + json: Mapping[str, object] | None, + timeout: float, + ) -> httpx.Request: ... + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: ... + + +def default_kill_switch_http_client() -> KillSwitchHttpClient: + return get_async_httpx_client(llm_provider=httpxSpecialProvider.AgentKillSwitch).client + + +KillSwitchAuditLogWriter: TypeAlias = Callable[[LiteLLM_AuditLogs], Awaitable[None]] # mutable-ok: Callable params + + +def default_kill_switch_audit_log_writer() -> KillSwitchAuditLogWriter: + return create_audit_log_for_update + + +def build_kill_switch_audit_log( + *, + result: AgentKillSwitchResult, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str | None, +) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.AGENT_TABLE_NAME, + object_id=result.agent_id, + action="kill_switch_fired", + updated_values=result.model_dump_json(exclude_none=True), + ) + + +async def fire_kill_switch( + *, + agent_id: str, + config: AgentKillSwitchConfig, + http_client: KillSwitchHttpClient, + timeout: float = AGENT_KILL_SWITCH_TIMEOUT_SECONDS, +) -> AgentKillSwitchResult: + request: Final = build_kill_switch_request(config) + reported_url: Final = str(httpx.URL(request.url).copy_with(query=None)) + verbose_proxy_logger.info("Firing kill switch for agent %s: %s %s", agent_id, request.method, reported_url) + try: + response: Final = await http_client.send( + http_client.build_request( + request.method, + request.url, + headers=request.headers, + json=request.json_body, + timeout=timeout, + ), + stream=True, + follow_redirects=False, + ) + body: Final = await _read_text_prefix(response, AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS) + except httpx.HTTPError as exc: + verbose_proxy_logger.warning("Kill switch for agent %s failed: %s", agent_id, type(exc).__name__) + return AgentKillSwitchResult( + agent_id=agent_id, + url=reported_url, + method=config.method, + error=type(exc).__name__, + ) + return AgentKillSwitchResult( + agent_id=agent_id, + url=reported_url, + method=config.method, + status_code=response.status_code, + response_body=body, + ) + + +async def _read_text_prefix(response: httpx.Response, max_chars: int) -> str: + try: + return await _take_text(response.aiter_text(), max_chars) + finally: + await response.aclose() + + +async def _take_text(chunks: AsyncIterator[str], max_chars: int) -> str: + taken = "" # rebind-ok: running prefix of a stream that is abandoned once the cap is hit + async for chunk in chunks: + taken += chunk # rebind-ok: see above + if len(taken) >= max_chars: + break + return taken[:max_chars] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 7f8d8c6af66..f7aef09fa29 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,8 +1,9 @@ from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias +from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, PrivateAttr, StrictInt +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field_validator from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -178,6 +179,74 @@ class AgentObjectPermission(TypedDict, total=False): agents: list[str] | None +class AgentKillSwitchBearerAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["bearer"] + token: str + + +class AgentKillSwitchApiKeyAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["api_key"] + header_name: str = "x-api-key" + api_key: str + + +class AgentKillSwitchBasicAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["basic"] + username: str + password: str + + +AgentKillSwitchAuth: TypeAlias = Annotated[ + AgentKillSwitchBearerAuth | AgentKillSwitchApiKeyAuth | AgentKillSwitchBasicAuth, + Field(discriminator="type"), +] + +AgentKillSwitchMethod: TypeAlias = Literal["POST", "PUT", "PATCH", "DELETE", "GET"] + + +class AgentKillSwitchConfig(BaseModel): + """Webhook an admin fires to shut an agent down out of band. LiteLLM only + makes the call; whatever the endpoint does with it is the agent's business.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + url: str + method: AgentKillSwitchMethod = "POST" + headers: Mapping[str, str] = Field(default_factory=dict) + query_params: Mapping[str, str] = Field(default_factory=dict) + body: Mapping[str, object] | None = None + auth: AgentKillSwitchAuth | None = None + + @field_validator("url") + @classmethod + def _require_absolute_http_url(cls, value: str) -> str: + parts: Final = urlsplit(value) + if parts.scheme not in ("http", "https") or not parts.netloc: + raise ValueError("kill_switch.url must be an absolute http(s) URL") + return value + + +class AgentKillSwitchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + url: str + method: AgentKillSwitchMethod + status_code: int | None = None + response_body: str | None = None + error: str | None = None + + @property + def succeeded(self) -> bool: + return self.status_code is not None and 200 <= self.status_code < 300 + + class AgentConfig(TypedDict, total=False): agent_name: Required[str] agent_card_params: Required[AgentCard] @@ -190,6 +259,7 @@ class AgentConfig(TypedDict, total=False): static_headers: dict[str, str] | None extra_headers: list[str] | None access_group_ids: ReadOnly[Sequence[str] | None] + kill_switch: ReadOnly[AgentKillSwitchConfig | None] class PatchAgentRequest(TypedDict, total=False): @@ -204,6 +274,7 @@ class PatchAgentRequest(TypedDict, total=False): static_headers: dict[str, str] | None extra_headers: list[str] | None access_group_ids: ReadOnly[Sequence[str] | None] + kill_switch: ReadOnly[AgentKillSwitchConfig | None] AGENT_CALLER_USER_ID_HEADER: Final = "x-litellm-user-id" @@ -243,6 +314,7 @@ class AgentResponse(BaseModel): static_headers: dict[str, str] | None = None extra_headers: list[str] | None = None access_group_ids: Sequence[str] | None = None + kill_switch: AgentKillSwitchConfig | None = None keys: list[AgentKeySummary] | None = None search_score: float | None = None created_at: datetime | None = None diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 06982a16755..fa2d1373ea1 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -26,6 +26,7 @@ class httpxSpecialProvider(str, Enum): RAG = "rag" A2AProvider = "a2a_provider" AgentHealthCheck = "agent_health_check" + AgentKillSwitch = "agent_kill_switch" A2A = "a2a" PromptManagement = "prompt_management" UI = "ui" diff --git a/schema.prisma b/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index b036e0dac4d..ef20e88c368 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -451,7 +451,7 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( - return_value=SimpleNamespace(litellm_params={}, object_permission_id=None) + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=None) ) mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None) @@ -736,6 +736,7 @@ async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted(): "model": "bedrock/agentcore/my-agent", }, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -784,6 +785,7 @@ async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely(): return_value=SimpleNamespace( litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -830,6 +832,7 @@ async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_ } }, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -878,6 +881,7 @@ async def test_update_agent_in_db_clears_secret_on_explicit_empty_value(): return_value=SimpleNamespace( litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -1110,7 +1114,9 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( - return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"]) + return_value=SimpleNamespace( + litellm_params={}, object_permission_id=None, kill_switch=None, access_group_ids=["ag-1"] + ) ) mock_update = AsyncMock(return_value=_agent_row_mock(expected)) mock_prisma.db.litellm_agentstable.update = mock_update @@ -1126,3 +1132,155 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro ) assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) + + +_KILL_SWITCH: Final = { + "url": "https://ops.example.com/kill", + "method": "POST", + "headers": {"X-Env": "prod"}, + "query_params": {"reason": "manual"}, + "body": {"action": "stop"}, + "auth": {"type": "bearer", "token": "tok-real"}, +} + + +@pytest.mark.asyncio +async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unset(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "kill_switch": _KILL_SWITCH, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + assert json.loads(mock_create.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH + + await registry.add_agent_to_db( + agent={"agent_name": "Plain Agent", "agent_card_params": _sample_agent_card_params()}, + prisma_client=mock_prisma, + created_by="test-user", + ) + assert mock_create.call_args.kwargs["data"]["kill_switch"] == json.dumps(None) + + +@pytest.mark.asyncio +async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.create = AsyncMock(return_value=_agent_row_mock([])) + + with pytest.raises(Exception, match="absolute http"): + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "kill_switch": {**_KILL_SWITCH, "url": "ops.example.com/kill"}, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + mock_prisma.db.litellm_agentstable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on_null(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old", + "litellm_params": {}, + "object_permission_id": None, + "kill_switch": _KILL_SWITCH, + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"agent_name": "New"}, prisma_client=mock_prisma, updated_by="u" + ) + assert "kill_switch" not in mock_update.call_args.kwargs["data"] + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"kill_switch": None}, prisma_client=mock_prisma, updated_by="u" + ) + assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None), ( + "prisma-client-py silently drops None, so the clear must be written as the JSON literal null" + ) + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_the_marker(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "A", + "litellm_params": {}, + "object_permission_id": None, + "kill_switch": _KILL_SWITCH, + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={ + "kill_switch": { + **_KILL_SWITCH, + "url": "https://ops.example.com/v2/kill", + "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}, + } + }, + prisma_client=mock_prisma, + updated_by="u", + ) + + assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == { + **_KILL_SWITCH, + "url": "https://ops.example.com/v2/kill", + } + + +@pytest.mark.asyncio +async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_secret_when_echoed(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH)) + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + base: Final = {"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params(), "litellm_params": {}} + + await registry.update_agent_in_db(agent_id="agent-123", agent=base, prisma_client=mock_prisma, updated_by="u") + assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None) + + echoed: Final = {**_KILL_SWITCH, "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}} + await registry.update_agent_in_db( + agent_id="agent-123", agent={**base, "kill_switch": echoed}, prisma_client=mock_prisma, updated_by="u" + ) + assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH + + +def test_load_agents_from_config_exposes_a_typed_kill_switch(): + registry: Final = AgentRegistry() + + registry.load_agents_from_config( + [{"agent_name": "cfg-agent", "agent_card_params": _sample_agent_card_params(), "kill_switch": _KILL_SWITCH}] + ) + + (agent,) = registry.get_agent_list() + assert agent.kill_switch is not None + assert agent.kill_switch.model_dump() == _KILL_SWITCH diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 482294e7b92..526f24c5221 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1,13 +1,15 @@ import json +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.constants import REDACTED_BY_LITELM_STRING -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( RestrictedAgentAccess, @@ -1136,3 +1138,222 @@ def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch assert duplicate.status_code == 400 assert "already in public agent groups" in duplicate.json()["detail"] + + +_KILL_SWITCH: Final = { + "url": "https://ops.example.com/kill", + "method": "POST", + "headers": {"X-Env": "prod"}, + "query_params": {"reason": "manual"}, + "body": {"action": "stop"}, + "auth": {"type": "bearer", "token": "tok-real"}, +} + + +def _agent_with_kill_switch() -> AgentResponse: + return AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={}, + kill_switch=_KILL_SWITCH, + ) + + +class _FakeKillSwitchClient: + def __init__(self, response: httpx.Response) -> None: + self.calls: list[tuple[str, str, dict[str, str], object, float]] = [] # mutable-ok: test double records calls + self._response: Final = response + + def build_request(self, method: str, url: str, *, headers, json, timeout: float) -> httpx.Request: + self.calls.append((method, url, dict(headers), json, timeout)) + return httpx.Request(method, url, headers=dict(headers), json=json) + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: + return self._response + + +class _AuditLogRecorder: + def __init__(self) -> None: + self.rows: list[LiteLLM_AuditLogs] = [] # mutable-ok: test double records writes + + async def __call__(self, request_data: LiteLLM_AuditLogs) -> None: + self.rows.append(request_data) + + +def _kill_switch_app( + role: LitellmUserRoles, + http_client: _FakeKillSwitchClient, + audit_log: _AuditLogRecorder | None = None, +) -> TestClient: + test_client: Final = _make_app_with_role(role) + test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_http_client] = lambda: http_client + test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_audit_log_writer] = ( + lambda: audit_log or _AuditLogRecorder() + ) + return test_client + + +def test_kill_switch_trigger_fires_the_configured_webhook_and_returns_the_result(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200, text="ok")) + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 200, resp.text + assert resp.json() == { + "agent_id": "agent-123", + "url": "https://ops.example.com/kill", + "method": "POST", + "status_code": 200, + "response_body": "ok", + "error": None, + } + (method, url, headers, body, _timeout) = fake.calls[0] + assert (method, url, body) == ("POST", "https://ops.example.com/kill?reason=manual", {"action": "stop"}) + assert headers == {"X-Env": "prod", "Authorization": "Bearer tok-real"} + + +def test_kill_switch_trigger_writes_an_audit_log_row_naming_the_admin_and_the_sanitized_result(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(202, text='{"stopped": true}')) + audit: Final = _AuditLogRecorder() + test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit) + test_client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-k" + ) + + resp: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}) + + assert resp.status_code == 200, resp.text + (row,) = audit.rows + assert (row.action, row.table_name, row.object_id) == ( + "kill_switch_fired", + LitellmTableNames.AGENT_TABLE_NAME, + "agent-123", + ) + assert (row.changed_by, row.changed_by_api_key) == ("test-user", "hashed-k") + assert row.before_value is None + assert json.loads(row.updated_values) == { + "agent_id": "agent-123", + "url": "https://ops.example.com/kill", + "method": "POST", + "status_code": 202, + "response_body": '{"stopped": true}', + } + assert "tok-real" not in row.model_dump_json() + + +def test_kill_switch_trigger_returns_502_and_still_audits_when_the_webhook_rejects(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(401, text="bad token")) + audit: Final = _AuditLogRecorder() + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 502, resp.text + assert resp.json()["detail"]["status_code"] == 401 + assert resp.json()["detail"]["response_body"] == "bad token" + (row,) = audit.rows + assert row.action == "kill_switch_fired" + assert json.loads(row.updated_values)["status_code"] == 401 + + +@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +def test_kill_switch_trigger_is_refused_before_any_webhook_call_for_non_admins(monkeypatch, role) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200)) + audit: Final = _AuditLogRecorder() + + resp: Final = _kill_switch_app(role, fake, audit).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 403, resp.text + assert fake.calls == [] + assert audit.rows == [] + + +def test_kill_switch_trigger_404s_unknown_agent_and_400s_an_agent_without_one(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(side_effect=[None, _sample_agent_response()]) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200)) + audit: Final = _AuditLogRecorder() + test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit) + + missing: Final = test_client.post("/v1/agents/nope/kill_switch", headers={"Authorization": "Bearer k"}) + unconfigured: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}) + + assert missing.status_code == 404 + assert unconfigured.status_code == 400 + assert "no kill_switch configured" in unconfigured.json()["detail"] + assert fake.calls == [] + assert audit.rows == [] + + +def test_kill_switch_trigger_fires_the_db_row_config_over_a_stale_in_memory_copy(monkeypatch) -> None: + """Another replica may have updated the agent; the row is the source of truth for what gets fired.""" + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + db_row: Final = SimpleNamespace( + agent_id="agent-123", + kill_switch={"url": "https://ops.example.com/kill-v2", "method": "DELETE", "auth": None}, + ) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=db_row) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + fake: Final = _FakeKillSwitchClient(httpx.Response(204)) + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 200, resp.text + (method, url, headers, body, _timeout) = fake.calls[0] + assert (method, url, headers, body) == ("DELETE", "https://ops.example.com/kill-v2", {}, None) + assert prisma.db.litellm_agentstable.find_unique.await_args.kwargs == {"where": {"agent_id": "agent-123"}} + registry.get_agent_by_id.assert_not_called() + + +def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_others(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + registry.ids_for_agent = MagicMock(return_value=("agent-123",)) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + + def _get_as(role: LitellmUserRoles): + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"}) + + admin: Final = _get_as(LitellmUserRoles.PROXY_ADMIN) + assert admin.status_code == 200, admin.text + assert admin.json()["kill_switch"] == { + **_KILL_SWITCH, + "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}, + } + + internal: Final = _get_as(LitellmUserRoles.INTERNAL_USER) + assert internal.status_code == 200, internal.text + assert internal.json()["kill_switch"] is None + assert "tok-real" not in internal.text diff --git a/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py b/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py new file mode 100644 index 00000000000..a6bb945713e --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py @@ -0,0 +1,248 @@ +from base64 import b64encode +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.constants import REDACTED_BY_LITELM_STRING +from litellm.proxy.agent_endpoints.kill_switch import ( + build_kill_switch_request, + fire_kill_switch, + redact_kill_switch, + restore_kill_switch, +) +from litellm.types.agents import AgentKillSwitchConfig + + +@dataclass(frozen=True, slots=True) +class _SentRequest: + method: str + url: str + headers: Mapping[str, str] + json: Mapping[str, object] | None + timeout: float + + +class _RecordingClient: + def __init__(self, respond: httpx.Response | httpx.HTTPError) -> None: + self.sent: list[_SentRequest] = [] # mutable-ok: test double records calls + self.follow_redirects: list[bool] = [] # mutable-ok: test double records calls + self._respond: Final = respond + + def build_request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + json: Mapping[str, object] | None, + timeout: float, + ) -> httpx.Request: + self.sent.append(_SentRequest(method, url, headers, json, timeout)) + return httpx.Request(method, url, headers=dict(headers), json=json) + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: + self.follow_redirects.append(follow_redirects) + if isinstance(self._respond, httpx.HTTPError): + raise self._respond + return self._respond + + +class _CountingStream(httpx.AsyncByteStream): + def __init__(self, chunk: bytes, chunks: int) -> None: + self.pulled: int = 0 # rebind-ok: test double counts reads + self._chunk: Final = chunk + self._chunks: Final = chunks + + async def __aiter__(self): + for _ in range(self._chunks): + self.pulled += 1 # rebind-ok: test double counts reads + yield self._chunk + + +def _config(**overrides: object) -> AgentKillSwitchConfig: + return AgentKillSwitchConfig.model_validate({"url": "https://ops.example.com/agents/kill", **overrides}) + + +def test_request_carries_endpoint_method_query_params_headers_and_body() -> None: + request: Final = build_kill_switch_request( + _config( + url="https://ops.example.com/kill?env=prod", + method="PUT", + query_params={"agent": "billing-bot", "reason": "manual stop"}, + headers={"X-Trace": "abc"}, + body={"action": "stop", "hard": True}, + ) + ) + + assert request.method == "PUT" + assert str(httpx.URL(request.url)) == "https://ops.example.com/kill?env=prod&agent=billing-bot&reason=manual+stop" + assert dict(request.headers) == {"X-Trace": "abc"} + assert request.json_body == {"action": "stop", "hard": True} + + +def test_request_defaults_to_post_with_no_body_and_untouched_url() -> None: + request: Final = build_kill_switch_request(_config()) + + assert (request.method, request.url, dict(request.headers), request.json_body) == ( + "POST", + "https://ops.example.com/agents/kill", + {}, + None, + ) + + +@pytest.mark.parametrize( + ("auth", "expected_headers"), + [ + ({"type": "bearer", "token": "tok-123"}, {"Authorization": "Bearer tok-123"}), + ({"type": "api_key", "api_key": "k-456"}, {"x-api-key": "k-456"}), + ({"type": "api_key", "header_name": "X-Ops-Key", "api_key": "k-456"}, {"X-Ops-Key": "k-456"}), + ( + {"type": "basic", "username": "ops", "password": "pw:1"}, + {"Authorization": f"Basic {b64encode(b'ops:pw:1').decode()}"}, + ), + ], +) +def test_auth_becomes_the_matching_request_header(auth: Mapping[str, object], expected_headers: dict[str, str]) -> None: + request: Final = build_kill_switch_request(_config(auth=auth)) + + assert dict(request.headers) == expected_headers + + +def test_auth_header_wins_over_a_conflicting_custom_header() -> None: + request: Final = build_kill_switch_request( + _config(headers={"Authorization": "stale", "X-Env": "prod"}, auth={"type": "bearer", "token": "fresh"}) + ) + + assert dict(request.headers) == {"Authorization": "Bearer fresh", "X-Env": "prod"} + + +@pytest.mark.parametrize("url", ["ftp://ops.example.com/kill", "/relative/kill", "ops.example.com/kill", ""]) +def test_config_rejects_non_http_urls(url: str) -> None: + with pytest.raises(ValidationError, match="absolute http"): + _config(url=url) + + +def test_config_rejects_unknown_auth_type_and_unknown_fields() -> None: + with pytest.raises(ValidationError): + _config(auth={"type": "hmac", "secret": "x"}) + with pytest.raises(ValidationError): + _config(endpoint="https://typo.example.com") + + +@pytest.mark.parametrize( + ("auth", "secret_field"), + [ + ({"type": "bearer", "token": "tok-123"}, "token"), + ({"type": "api_key", "header_name": "X-K", "api_key": "k-456"}, "api_key"), + ({"type": "basic", "username": "ops", "password": "pw"}, "password"), + ], +) +def test_redact_replaces_only_the_secret_and_restore_puts_it_back(auth: dict[str, str], secret_field: str) -> None: + original: Final = _config(auth=auth) + + redacted: Final = redact_kill_switch(original) + assert redacted is not None and redacted.auth is not None + assert redacted.auth.model_dump() == {**auth, secret_field: REDACTED_BY_LITELM_STRING} + assert original.auth is not None and original.auth.model_dump() == auth, "redact must not mutate its input" + + restored: Final = restore_kill_switch(redacted, original) + assert restored == original + + +def test_restore_keeps_a_rotated_secret_and_never_stores_the_marker_itself() -> None: + rotated: Final = _config(auth={"type": "bearer", "token": "new-token"}) + stored: Final = _config(auth={"type": "bearer", "token": "old-token"}) + assert restore_kill_switch(rotated, stored) == rotated + assert restore_kill_switch(None, stored) is None + + marker_only: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING}) + assert restore_kill_switch(marker_only, None) == _config(auth={"type": "bearer", "token": ""}) + + +def test_restore_does_not_borrow_a_secret_from_a_different_auth_type() -> None: + incoming: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING}) + stored: Final = _config(auth={"type": "api_key", "api_key": "k-456"}) + + assert restore_kill_switch(incoming, stored) == _config(auth={"type": "bearer", "token": ""}) + + +def test_redact_passes_through_configs_without_auth() -> None: + assert redact_kill_switch(None) is None + plain: Final = _config(headers={"X-Env": "prod"}) + assert redact_kill_switch(plain) is plain + + +@pytest.mark.asyncio +async def test_fire_sends_exactly_the_built_request_and_reports_the_2xx_reply_without_the_query() -> None: + client: Final = _RecordingClient(httpx.Response(202, text="stopping")) + config: Final = _config( + method="DELETE", + query_params={"force": "1", "token": "qs-secret"}, + headers={"X-Env": "prod"}, + body={"agent": "billing-bot"}, + auth={"type": "bearer", "token": "tok-123"}, + ) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=config, http_client=client, timeout=3.5) + + assert client.sent == [ + _SentRequest( + method="DELETE", + url="https://ops.example.com/agents/kill?force=1&token=qs-secret", + headers={"X-Env": "prod", "Authorization": "Bearer tok-123"}, + json={"agent": "billing-bot"}, + timeout=3.5, + ) + ] + assert client.follow_redirects == [False], "a redirecting webhook must not be followed to another host" + assert result.succeeded is True + assert result.model_dump() == { + "agent_id": "agent-1", + "url": "https://ops.example.com/agents/kill", + "method": "DELETE", + "status_code": 202, + "response_body": "stopping", + "error": None, + } + + +@pytest.mark.asyncio +async def test_fire_reports_a_non_2xx_reply_as_failure_with_the_body() -> None: + client: Final = _RecordingClient(httpx.Response(503, text="x" * 5000)) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client) + + assert result.succeeded is False + assert result.status_code == 503 + assert result.response_body == "x" * 2000 + assert result.error is None + + +@pytest.mark.asyncio +async def test_fire_stops_reading_the_body_at_the_cap_instead_of_buffering_the_whole_reply() -> None: + stream: Final = _CountingStream(b"y" * 500, chunks=100) + client: Final = _RecordingClient(httpx.Response(200, stream=stream)) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client) + + assert result.response_body == "y" * 2000 + assert stream.pulled == 4, f"read {stream.pulled} of 100 chunks for a 2000 char cap" + + +@pytest.mark.asyncio +async def test_fire_reports_a_transport_error_by_type_without_raising_or_echoing_the_url() -> None: + client: Final = _RecordingClient(httpx.ConnectError("boom https://ops.example.com/agents/kill?token=qs-secret")) + + result: Final = await fire_kill_switch( + agent_id="agent-1", config=_config(query_params={"token": "qs-secret"}), http_client=client + ) + + assert result.succeeded is False + assert (result.status_code, result.response_body) == (None, None) + assert result.error == "ConnectError" + assert "qs-secret" not in result.model_dump_json() diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 7bb79a115dd..f76a02e8361 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3776,6 +3776,7 @@ AGENT_MANAGEMENT_ROUTES = [ "/v1/agents/abc-123", "/v1/agents/make_public", "/v1/agents/abc-123/make_public", + "/v1/agents/abc-123/kill_switch", ] AGENT_INFERENCE_ROUTES = [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index 8e100d0c3ed..3d863036234 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -27,6 +27,7 @@ import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/component import { Input } from "@/components/ui/input"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { Field, FieldDescription, FieldError, FieldGroup, FieldLabel } from "@/components/ui/field"; +import type { KeyValueFormValue, KillSwitchConfig, KillSwitchFormValue } from "./kill_switch_config"; export interface AgentSkillFormValue { id?: string; @@ -54,6 +55,8 @@ export type AgentFormFieldValue = | string[] | AgentSkillFormValue[] | StaticHeaderFormValue[] + | KeyValueFormValue[] + | KillSwitchFormValue | McpServerSelection | Record | null @@ -82,6 +85,7 @@ export interface AgentFormValues { output_cost_per_token?: string | number; static_headers?: StaticHeaderFormValue[]; extra_headers?: string[]; + kill_switch?: KillSwitchFormValue; tpm_limit?: number | null; rpm_limit?: number | null; session_tpm_limit?: number | null; @@ -123,6 +127,7 @@ export interface AgentRequestPayload { litellm_params?: Record; object_permission?: Record; access_group_ids?: string[]; + kill_switch?: KillSwitchConfig | null; } interface AgentFormFieldProps { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx new file mode 100644 index 00000000000..aa4b83478b6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx @@ -0,0 +1,124 @@ +import React from "react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import AgentKillSwitchDangerZone from "./AgentKillSwitchDangerZone"; +import * as networking from "@/components/networking"; +import { toast } from "@/lib/toast"; + +vi.mock("@/components/networking", () => ({ + triggerAgentKillSwitchCall: vi.fn(), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), error: vi.fn() }, +})); + +const killSwitch = { url: "https://ops.example.com/kill", method: "DELETE" as const }; + +const renderZone = (props: Partial> = {}) => + render( + , + ); + +const openDialog = () => { + fireEvent.click(screen.getByRole("button", { name: "Fire Kill Switch" })); + return screen.getByRole("dialog"); +}; + +const dialogFireButton = () => within(screen.getByRole("dialog")).getByRole("button", { name: "Fire Kill Switch" }); + +describe("AgentKillSwitchDangerZone", () => { + beforeEach(() => { + vi.mocked(networking.triggerAgentKillSwitchCall).mockReset(); + vi.mocked(toast.success).mockReset(); + vi.mocked(toast.error).mockReset(); + }); + + it("renders nothing for non-admins", () => { + const { container } = renderZone({ isAdmin: false }); + + expect(container).toBeEmptyDOMElement(); + }); + + it("shows the webhook target and an outage warning inside a Danger Zone region", () => { + renderZone(); + + const region = screen.getByRole("region", { name: "Danger Zone" }); + expect(region).toHaveTextContent("DELETE https://ops.example.com/kill"); + expect(region).toHaveTextContent("can cause an outage"); + expect(screen.getByRole("button", { name: "Fire Kill Switch" })).toBeEnabled(); + }); + + it("shows an unconfigured notice without a fire button when no kill switch is set", () => { + renderZone({ killSwitch: null }); + + expect(screen.getByRole("region", { name: "Danger Zone" })).toHaveTextContent("Not configured"); + expect(screen.queryByRole("button", { name: "Fire Kill Switch" })).not.toBeInTheDocument(); + }); + + it("keeps the confirm button disabled until the exact agent name is typed", () => { + renderZone(); + openDialog(); + + expect(dialogFireButton()).toBeDisabled(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agen" } }); + expect(dialogFireButton()).toBeDisabled(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + expect(dialogFireButton()).toBeEnabled(); + expect(networking.triggerAgentKillSwitchCall).not.toHaveBeenCalled(); + }); + + it("fires the webhook after typed confirmation, closes the dialog and shows the sanitized result", async () => { + const firedResult = { + agent_id: "agent-1", + url: killSwitch.url, + method: "DELETE" as const, + status_code: 202, + response_body: '{"stopped": true}', + }; + vi.mocked(networking.triggerAgentKillSwitchCall).mockResolvedValue(firedResult); + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(dialogFireButton()); + + expect(await screen.findByRole("status")).toHaveTextContent('Last result: HTTP 202 {"stopped": true}'); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(networking.triggerAgentKillSwitchCall).toHaveBeenCalledWith("sk-test", "agent-1"); + expect(toast.success).toHaveBeenCalledWith("Kill switch fired (HTTP 202)"); + }); + + it("does not call the webhook when the dialog is cancelled", () => { + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(screen.getByRole("button", { name: "Cancel" })); + + expect(networking.triggerAgentKillSwitchCall).not.toHaveBeenCalled(); + expect(screen.queryByRole("status")).not.toBeInTheDocument(); + }); + + it("surfaces a failed webhook as an error toast and keeps the dialog open", async () => { + vi.mocked(networking.triggerAgentKillSwitchCall).mockRejectedValue(new Error("Kill switch webhook returned 500")); + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(dialogFireButton()); + + await waitFor(() => expect(toast.error).toHaveBeenCalledWith("Kill switch webhook returned 500")); + expect(screen.getByRole("dialog")).toBeInTheDocument(); + expect(screen.queryByRole("status")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx new file mode 100644 index 00000000000..8d68a8f921b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx @@ -0,0 +1,147 @@ +import { CircleAlert } from "lucide-react"; +import React, { useState } from "react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { toast } from "@/lib/toast"; +import { AgentKillSwitchResult, triggerAgentKillSwitchCall } from "@/components/networking"; +import { KillSwitchConfig } from "./kill_switch_config"; + +interface AgentKillSwitchDangerZoneProps { + agentId: string; + agentName: string; + killSwitch: KillSwitchConfig | null | undefined; + accessToken: string | null; + isAdmin: boolean; +} + +const AgentKillSwitchDangerZone: React.FC = ({ + agentId, + agentName, + killSwitch, + accessToken, + isAdmin, +}) => { + const [isConfirmOpen, setIsConfirmOpen] = useState(false); + const [confirmationInput, setConfirmationInput] = useState(""); + const [isFiring, setIsFiring] = useState(false); + const [lastResult, setLastResult] = useState(null); + + if (!isAdmin) return null; + + const openConfirm = () => { + setConfirmationInput(""); + setIsConfirmOpen(true); + }; + + const fire = async () => { + if (!accessToken) return; + setIsFiring(true); + setLastResult(null); + try { + const result = await triggerAgentKillSwitchCall(accessToken, agentId); + setLastResult(result); + setIsConfirmOpen(false); + toast.success(`Kill switch fired (HTTP ${result.status_code})`); + } catch (error) { + toast.error(error instanceof Error ? error.message : "Failed to fire kill switch"); + } finally { + setIsFiring(false); + } + }; + + return ( +
+

+ Danger Zone +

+
+
+
+

Kill switch

+ {killSwitch ? ( + <> +

+ Calls the configured webhook to stop this agent's upstream runtime. This can cause an outage for + everyone using the agent and cannot be undone from LiteLLM +

+

+ {killSwitch.method ?? "POST"} {killSwitch.url} +

+ + ) : ( +

+ Not configured. Add a kill switch webhook under Settings to enable this action +

+ )} +
+ {killSwitch && ( + + )} +
+ {lastResult && ( +

+ Last result: HTTP {lastResult.status_code} + {lastResult.response_body ? ` ${lastResult.response_body}` : ""} +

+ )} +
+ + !open && !isFiring && setIsConfirmOpen(false)}> + + + Fire kill switch for {agentName}? + +
+ + + This can cause an outage + + LiteLLM will call {killSwitch?.method ?? "POST"} {killSwitch?.url} immediately. Whatever that webhook + does to the agent is outside LiteLLM's control and cannot be reverted here + + +
+

+ Type {agentName} to confirm: +

+ + + + + setConfirmationInput(e.target.value)} + placeholder={agentName} + autoFocus + /> + +
+
+ + + + + + +
+ ); +}; + +export default AgentKillSwitchDangerZone; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx new file mode 100644 index 00000000000..578d80da6ed --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx @@ -0,0 +1,223 @@ +import React from "react"; +import { useFieldArray, useFormContext, useWatch } from "react-hook-form"; +import { Plus, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; +import { Field, FieldTitle } from "@/components/ui/field"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { AgentFormField, AgentFormValues, labelWithHint } from "./AgentFormKit"; +import { KILL_SWITCH_AUTH_TYPES, KILL_SWITCH_METHODS, validateKillSwitchBody } from "./kill_switch_config"; + +const KeyValueFieldArray = ({ + name, + addLabel, + keyPlaceholder, + valuePlaceholder, +}: { + name: "kill_switch.headers" | "kill_switch.query_params"; + addLabel: string; + keyPlaceholder: string; + valuePlaceholder: string; +}) => { + const { control } = useFormContext(); + const { fields, append, remove } = useFieldArray({ control, name }); + + return ( +
+ {fields.map((item, index) => ( +
+ + {({ value, onChange, ref, ...control }) => ( + + )} + + + {({ value, onChange, ref, ...control }) => ( + + )} + + +
+ ))} + +
+ ); +}; + +const TextField = ({ + name, + label, + placeholder, + required, + secret, +}: { + name: `kill_switch.${string}`; + label: React.ReactNode; + placeholder?: string; + required?: string; + secret?: boolean; +}) => ( + + {({ value, onChange, ref, ...control }) => + secret ? ( + + ) : ( + + ) + } + +); + +const KillSwitchAuthFields = () => { + const { control } = useFormContext(); + const authType = useWatch({ control, name: "kill_switch.auth_type" }); + + switch (authType) { + case "bearer": + return ; + case "api_key": + return ( + <> + + + + ); + case "basic": + return ( + <> + + + + ); + default: + return null; + } +}; + +const KillSwitchFormFields = () => ( + <> + + + + {({ value, onChange, ref: _ref, ...control }) => ( + + )} + + + + Headers + + + + + Query Parameters + + + + validateKillSwitchBody(typeof value === "string" ? value : "") }} + > + {({ value, onChange, ref, ...control }) => ( +