diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index fbdbfafa66d..b442948c3ca 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)") DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*") COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE) GLOB_CHARS = frozenset("*?") +DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"}) # Trees whose jobs are sharded with no catch-all bucket, so every child that holds # tests has to be named by some shard or it runs nowhere. A child listed here is @@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]: return frozenset(ast.literal_eval(values)) +def _owns(workflow: str, relative_path: str) -> bool: + owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow) + return owned_root is not None and relative_path.startswith(owned_root) + + +def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]: + return frozenset[str]().union( + *(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path)) + ) + + def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]: runner: Final = repo_root / "tests/integration/run.py" if not runner.exists(): @@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens ) != 1 ) - gha_tokens: Final = _invoked_test_tokens( - scalar - for path in (repo_root / ".github/workflows").glob("*.y*ml") - for scalar in _scalars(yaml.safe_load(path.read_text()), path.name) + workflow_tokens: Final = MappingProxyType( + { + path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name)) + for path in (repo_root / ".github/workflows").glob("*.y*ml") + } + ) + gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values()) + circle_owned: Final = ( + frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files ) findings: Final = ( tuple( - Finding(path, "integration contract is also selected by GitHub Actions") - for path in paths + Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions") + for path in sorted(circle_owned) if any(_token_covers(token, path) for token in gha_tokens) ) + tuple( @@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens for path in sorted(github_files) if not (repo_root / path).is_file() ) + + tuple( + Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow") + for path in sorted(github_files) + if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens)) + ) ) browser_commands: Final = tuple( scalar.value @@ -688,13 +710,9 @@ def main() -> int: scalars = _all_scalars() integration_paths, ownership_findings = _integration_ownership() - test_findings = ( - _uncovered_tests( - allowlist, - _invoked_selections(scalars) - + (Selection(included=integration_paths, ignored=frozenset()),), - ) - + ownership_findings + test_findings = _uncovered_tests( + allowlist, + _invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),), ) dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars)) stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles()) @@ -711,13 +729,25 @@ def main() -> int: test_findings, "Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.", ) + if ownership_findings: + browser_remedy = ( + " CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may " + "select them and the default Playwright config must exclude them." + if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings) + else "" + ) + _report( + "integration contracts with the wrong CI owner", + ownership_findings, + "Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy, + ) if dockerfile_findings: _report( "Dockerfiles that no CI job builds", dockerfile_findings, "Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.", ) - if stale_findings or test_findings or dockerfile_findings: + if stale_findings or test_findings or ownership_findings or dockerfile_findings: return 1 _write( diff --git a/AGENTS.md b/AGENTS.md index 87e4df49299..f4b1de57807 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a `tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones -End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md` +End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` @@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests, Always pull before starting any work. The checkout or worktree may be sitting on a stale branch -If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names +Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch diff --git a/litellm/__init__.py b/litellm/__init__.py index 88e56c14187..df52808643d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1177,7 +1177,9 @@ def _build_models_by_provider() -> dict: "edenai": edenai_models, "datarobot": datarobot_models, "vertex_ai": vertex_chat_models + | vertex_code_chat_models | vertex_text_models + | vertex_code_text_models | vertex_anthropic_models | vertex_vision_models | vertex_language_models @@ -1185,7 +1187,7 @@ def _build_models_by_provider() -> dict: | vertex_minimax_models | vertex_moonshot_models | vertex_zai_models, - "ai21": ai21_models, + "ai21": ai21_models | ai21_chat_models, "bedrock": bedrock_models | bedrock_converse_models, "petals": petals_models, "ollama": ollama_models, @@ -1416,6 +1418,7 @@ from .exceptions import ( InvalidRequestError, BadRequestError, ImageFetchError, + MalformedRequestFieldError as MalformedRequestFieldError, VectorStoreSearchError, NotFoundError, PermissionDeniedError, diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index ba542bebaf8..a3428261636 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -129,6 +129,7 @@ def get_assistants( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, aget_assistants=aget_assistants, ) @@ -312,6 +313,7 @@ def create_assistants( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, create_assistant_data=create_assistant_data, client=client, async_create_assistants=async_create_assistants, @@ -461,6 +463,7 @@ def delete_assistant( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, assistant_id=assistant_id, client=client, async_delete_assistants=async_delete_assistants, @@ -637,6 +640,7 @@ def create_thread( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, acreate_thread=acreate_thread, ) @@ -783,6 +787,7 @@ def get_thread( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, aget_thread=aget_thread, ) @@ -960,6 +965,7 @@ def add_message( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, a_add_message=a_add_message, ) @@ -1111,6 +1117,7 @@ def get_messages( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, aget_messages=aget_messages, ) @@ -1308,6 +1315,7 @@ def run_thread( timeout=timeout, max_retries=optional_params.max_retries, organization=organization, + litellm_params=litellm_params_dict, client=client, arun_thread=arun_thread, event_handler=event_handler, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index e7b009eaefe..a50cdf75d33 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -19,7 +19,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( ) from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage -from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS +from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS from litellm.utils import token_counter @@ -569,6 +569,7 @@ def extract_file_access_credentials( # A federated deployment holds no api_key, so without these the fetch that reads a # finished batch's output has nothing to authenticate with and its cost is never billed. *sorted(ANTHROPIC_WIF_KWARGS_KEYS), + *sorted(OPENAI_WIF_KWARGS_KEYS), ) for key in credential_keys: if key in litellm_params: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 1acd7ad67d1..9b5630d81ad 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -309,6 +309,7 @@ def create_batch( timeout=timeout, max_retries=optional_params.max_retries, _is_async=_is_async, + litellm_params=litellm_params, ) elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -471,6 +472,7 @@ def _handle_retrieve_batch_providers_without_provider_config( organization=organization, timeout=timeout, max_retries=optional_params.max_retries, + litellm_params=litellm_params, ) elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -819,6 +821,7 @@ def list_batches( organization=organization, timeout=timeout, max_retries=optional_params.max_retries, + litellm_params=litellm_params, ) elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -1016,6 +1019,7 @@ def cancel_batch( organization=organization, timeout=timeout, max_retries=optional_params.max_retries, + litellm_params=litellm_params, ) elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 979caa29832..8d23526b41a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -197,7 +197,7 @@ def _drops_audio_input(model: str, litellm_params: Mapping[str, object]) -> bool return not _supports_audio_input(model, litellm_params) -def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]: +def responses_provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]: return MappingProxyType( { key: value @@ -1455,7 +1455,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) model_response.id = _upstream_response_id(raw_response.id) or raw_response.id - for key, value in _provider_metadata(raw_response.model_extra).items(): + for key, value in responses_provider_metadata(raw_response.model_extra).items(): setattr(model_response, key, value) # Preserve hidden params from the ResponsesAPIResponse, especially the headers @@ -2139,7 +2139,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): from litellm.responses.utils import ResponseAPILoggingUtils usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_data.get("usage")) - provider_metadata: Final = _provider_metadata(response_data) + provider_metadata: Final = responses_provider_metadata(response_data) served_service_tier: Final = response_data.get("service_tier") return ModelResponseStream( choices=[ diff --git a/litellm/constants.py b/litellm/constants.py index 39363cc02e4..425fc22d340 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1859,6 +1859,7 @@ SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_L PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25)) PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0)) PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0)) +TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS: Final = 30 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500)) diff --git a/litellm/containers/README.md b/litellm/containers/README.md index 571bcaf415c..40723add854 100644 --- a/litellm/containers/README.md +++ b/litellm/containers/README.md @@ -120,6 +120,7 @@ class AzureContainerConfig(BaseContainerConfig): self, headers: dict, api_key: Optional[str] = None, + litellm_params: Optional[GenericLiteLLMParams] = None, ) -> dict: """Azure uses api-key header instead of Bearer token.""" import litellm diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0fb7363cbb5..5db809e486b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -86,6 +86,9 @@ from litellm.llms.openai.cost_calculation import ( from litellm.llms.openai.cost_calculation import ( video_output_cost_per_second, ) +from litellm.llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTAudioTranscriptionConfig, +) from litellm.llms.perplexity.cost_calculator import ( cost_per_token as perplexity_cost_per_token, ) @@ -325,6 +328,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool: + if model is None or custom_llm_provider != LlmProviders.OPENAI.value: + return False + transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI + ) + return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) + + OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] @@ -1739,6 +1751,12 @@ def completion_cost( audio_transcription_file_duration = _hidden.get( "audio_transcription_duration", getattr(completion_response, "duration", 0.0), + ) or ( + completion_response.usage.seconds + if isinstance(completion_response, TranscriptionResponse) + and isinstance(completion_response.usage, TranscriptionUsageDurationObject) + and _bills_provider_transcription_seconds(requested_model, custom_llm_provider) + else 0.0 ) elif call_type in _RERANK_CALL_TYPES: if completion_response is not None and isinstance(completion_response, RerankResponse): diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index ef35b61f10e..040d80e1ac0 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -458,7 +458,7 @@ def prepare_call( ) from error try: - resolved_api_base, resolved_api_key = provider_config.resolve_connection( + resolved_api_base, static_api_key = provider_config.resolve_connection( dynamic_api_base or api_base, dynamic_api_key or api_key ) except ValueError as error: @@ -469,6 +469,7 @@ def prepare_call( model=model, llm_provider=provider, ) + resolved_api_key: Final = provider_config.resolve_credential(static_api_key, resolved_api_base, kwargs) if resolved_api_key is None and provider_config.api_key_required: raise litellm.AuthenticationError( message=f"Missing API key for Decisions provider '{provider}'", diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 553fe4b884a..7d2ee03da0e 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -1,5 +1,7 @@ +import asyncio import warnings from collections.abc import Mapping +from functools import partial from typing import Final import httpx @@ -104,19 +106,22 @@ async def adecisions( videos: DecisionsVideoClips | None = None, **kwargs: object, ) -> OpenAIDecisionResponse | DecisionsResponse: - call: Final = _prepare_decisions_call( - model=model, - decision_input=input, - questions=questions, - safety_identifier=safety_identifier, - audio=audio, - videos=videos, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, + call: Final = await asyncio.to_thread( + partial( + _prepare_decisions_call, + model=model, + decision_input=input, + questions=questions, + safety_identifier=safety_identifier, + audio=audio, + videos=videos, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) ) return await asend(call) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index b3a0802db13..733f1aa72d5 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError): ) +class MalformedRequestFieldError(BadRequestError): + pass + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, diff --git a/litellm/files/main.py b/litellm/files/main.py index e259d4230de..16ff5b02f9c 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -259,6 +259,7 @@ def create_file( max_retries=optional_params.max_retries, organization=openai_creds.organization, create_file_data=_create_file_request, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "azure": azure_creds: Final = get_azure_credentials( @@ -388,6 +389,7 @@ def file_retrieve( timeout=timeout, max_retries=optional_params.max_retries, organization=openai_creds.organization, + litellm_params=optional_params.model_dump(exclude_none=True), ) elif custom_llm_provider == "azure": azure_creds: Final = get_azure_credentials( @@ -569,6 +571,7 @@ def file_delete( timeout=timeout, max_retries=optional_params.max_retries, organization=openai_creds.organization, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "azure": azure_creds: Final = get_azure_credentials( @@ -775,6 +778,7 @@ def file_list( timeout=timeout, max_retries=optional_params.max_retries, organization=openai_creds.organization, + litellm_params=optional_params.model_dump(exclude_none=True), ) elif custom_llm_provider == "azure": azure_creds: Final = get_azure_credentials( @@ -972,6 +976,7 @@ def file_content( timeout=timeout, max_retries=optional_params.max_retries, organization=openai_creds.organization, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "azure": azure_creds: Final = get_azure_credentials( @@ -1093,6 +1098,7 @@ def file_content_streaming( organization=openai_creds.organization, chunk_size=chunk_size, client=client if isinstance(client, (OpenAI, AsyncOpenAI)) else None, + litellm_params=optional_params.model_dump(exclude_none=True), ) elif custom_llm_provider == LlmProviders.VERTEX_AI.value: if not _is_async: diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 38be0666008..5297a8215be 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -239,6 +239,7 @@ def create_fine_tuning_job( client=kwargs.get( "client", None ), # note, when we add this to `GenericLiteLLMParams` it impacts a lot of other tests + linting + litellm_params=optional_params.model_dump(exclude_none=True), ) # Azure OpenAI elif custom_llm_provider == "azure": @@ -436,6 +437,7 @@ def cancel_fine_tuning_job( max_retries=optional_params.max_retries, _is_async=_is_async, client=kwargs.get("client", None), + litellm_params=optional_params.model_dump(exclude_none=True), ) # Azure OpenAI elif custom_llm_provider == "azure": @@ -591,6 +593,7 @@ def list_fine_tuning_jobs( max_retries=optional_params.max_retries, _is_async=_is_async, client=kwargs.get("client", None), + litellm_params=optional_params.model_dump(exclude_none=True), ) # Azure OpenAI elif custom_llm_provider == "azure": @@ -730,6 +733,7 @@ def retrieve_fine_tuning_job( max_retries=optional_params.max_retries, _is_async=_is_async, client=kwargs.get("client", None), + litellm_params=optional_params.model_dump(exclude_none=True), ) # Azure OpenAI elif custom_llm_provider == "azure": diff --git a/litellm/images/main.py b/litellm/images/main.py index be32ea4daed..7f94835d59b 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -39,6 +39,7 @@ from openai.types.audio.transcription_create_params import FileTypes # BFL handlers from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation +from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.main import ( azure_chat_completions, base_llm_aiohttp_handler, @@ -505,6 +506,7 @@ def image_generation( aimg_generation=aimg_generation, client=client, headers=headers, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "bedrock": if model is None: @@ -681,10 +683,13 @@ def image_variation( api_base = provider_config.get_api_base(litellm_params.get("api_base", None)) if image_variation_provider == LITELLM_IMAGE_VARIATION_PROVIDERS.OPENAI: - if api_key is None: - raise ValueError("API key is required for OpenAI image variations") if api_base is None: raise ValueError("API base is required for OpenAI image variations") + if api_key is None and ( + resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base, litellm_params=litellm_params) + is None + ): + raise ValueError("API key is required for OpenAI image variations") response = openai_image_variations.image_variations( model_response=model_response, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fa5496004e3..97a1073d7c0 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -27,6 +27,7 @@ from litellm.constants import ( PROXY_REJECTED_BEFORE_ROUTING_KEY, ) from litellm.exceptions import ( + MalformedRequestFieldError, validate_rate_limit_category, validate_rate_limit_type, ) @@ -3104,7 +3105,11 @@ class PrometheusLogger(CustomLogger): # On LiteLLM-side rejects (no deployment picked), route request_kwargs["model"] # into requested_model and leave deployment-scoped labels empty. - deployment_selected: Final = bool(model_id) and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY) + deployment_selected: Final = ( + bool(model_id) + and not _litellm_params.get(PROXY_REJECTED_BEFORE_ROUTING_KEY) + and not isinstance(exception, MalformedRequestFieldError) + ) if deployment_selected: label_litellm_model_name = litellm_model_name label_model_id = model_id diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 04755f9e1e9..9dee3714c60 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -78,6 +78,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "github_copilot_auth_type", "github_copilot_user_session", "fireworks_forward_user_id", + "responses_api_version", PROVIDER_AFFINITY_HEADER_KWARG_KEY, } ) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index fee48cca0cf..6e3769260cc 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -17,6 +17,7 @@ import random import sys import threading import time +import urllib.parse from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from datetime import datetime, timezone @@ -24,11 +25,11 @@ from importlib.resources import files from itertools import chain from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol +from typing import Final, Protocol, cast # noqa: TID251 # adapts boto3 and botocore's dynamic client responses import httpx from pydantic import ConfigDict, TypeAdapter -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm import verbose_logger from litellm.constants import ( @@ -266,6 +267,19 @@ class _SyncGetClient(Protocol): def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ... +class _S3Body(Protocol): + def read(self) -> bytes: ... + + +class _S3GetObjectResponse(TypedDict): + Body: ReadOnly[_S3Body] + ETag: NotRequired[ReadOnly[str]] + + +class _S3Client(Protocol): + def get_object(self, *, Bucket: str, Key: str) -> _S3GetObjectResponse: ... + + _FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable @@ -276,6 +290,20 @@ def _default_reload_client() -> _AsyncGetClient: return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap) +def _default_s3_client(timeout: int) -> _S3Client: + import boto3 + from botocore.config import Config + + return cast( # cast-ok: boto3's generic client factory does not expose its generated S3 response type + _S3Client, + boto3.client( # pyright: ignore[reportUnknownMemberType] # generic factory returns Unknown service clients + "s3", + region_name=os.getenv("AWS_REGION_NAME") or None, + config=Config(connect_timeout=timeout, read_timeout=timeout, retries={"total_max_attempts": 1}), + ), + ) + + def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome: reason: Final = f"{type(error).__name__} fetching {url}: {error}" if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)): @@ -283,7 +311,41 @@ def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) - return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None) +def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _FetchAttemptOutcome: + try: + parsed: Final[object] = json.loads(content) + except ValueError as e: + return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") + if not isinstance(parsed, dict): + return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}") + model_cost_map: Final = cast( # cast-ok: JSON object keys are strings by definition + dict[str, object], parsed + ) + return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag) + + +def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: + from botocore.exceptions import BotoCoreError, ClientError + + parts: Final = urllib.parse.urlsplit(url) + bucket: Final = parts.netloc + key: Final = parts.path.lstrip("/") + if not bucket or not key: + return ModelCostMapReloadUnavailable(reason=f"invalid S3 URL {url}: expected s3://bucket/key") + + try: + client: Final = _default_s3_client(timeout) + response: Final = client.get_object(Bucket=bucket, Key=key) + content: Final = response["Body"].read() + except (BotoCoreError, ClientError) as error: + return ModelCostMapReloadUnavailable(reason=f"S3 {type(error).__name__} fetching {url}: {error}") + return _classify_fetch_content(content, url, response.get("ETag")) + + async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: + parts: Final = urllib.parse.urlsplit(url) + if parts.scheme == "s3": + return await asyncio.to_thread(_attempt_fetch_s3, url, timeout) try: response: Final = await client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -292,6 +354,9 @@ async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _Fe def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: + parts: Final = urllib.parse.urlsplit(url) + if parts.scheme == "s3": + return _attempt_fetch_s3(url, timeout) try: response: Final = client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -307,15 +372,7 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp ) if response.is_error: return ModelCostMapReloadUnavailable(reason=f"HTTP {response.status_code} from {url}") - try: - parsed: Final = response.json() - except ValueError as e: - return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") - if not isinstance(parsed, dict): - return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}") - return ModelCostMapReloaded( - model_cost_map=parsed, revision=git_blob_id(response.content), etag=response.headers.get("etag") - ) + return _classify_fetch_content(response.content, url, response.headers.get("etag")) def _next_retry_wait( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 198b66a0268..8c310011556 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2311,6 +2311,13 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def attempt_scoped_copy(self) -> "Logging": + """A copy with its own model_call_details, so logging one attempt of a request through it leaves + the state the request's other attempts share untouched.""" + attempt: Final = copy.copy(self) + attempt.model_call_details = dict(self.model_call_details) + return attempt + def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" @@ -4357,6 +4364,9 @@ class Logging(LiteLLMLoggingBaseClass): row is built from this response rather than re-parsed from the served body.""" self.bridged_model_response = response + def record_bridged_responses_api_response(self, response: ResponsesAPIResponse) -> None: + self.record_bridged_model_response(self._translate_responses_api_response_to_model_response(response)) + def billable_result(self, result: object) -> object: """The response the cost calculator prices: a /v1/messages body served through the chat bridge is replaced by the ModelResponse its spend row is built from, anything else is @@ -4457,6 +4467,7 @@ class Logging(LiteLLMLoggingBaseClass): """ from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, + responses_provider_metadata, ) served_id: Final = _provider_response_id(result) @@ -4481,7 +4492,9 @@ class Logging(LiteLLMLoggingBaseClass): "usage-only ModelResponse to keep the spend_logs row.", str(e), ) - model_response: Final = litellm.ModelResponse(id=served_id) + model_response: Final = litellm.ModelResponse( + id=served_id, **responses_provider_metadata(result.model_extra) + ) model_response.model = self.model usage: Final = result.usage if usage is not None and ResponseAPILoggingUtils.is_response_api_usage(usage): diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index ce5dceb2e5c..04510c20db0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -858,9 +858,9 @@ class StandardBuiltInToolCostTracking: @staticmethod def get_web_search_options(kwargs: Mapping[str, object]) -> WebSearchOptions | None: - if "web_search_options" in kwargs: - web_search_options: Final = cast(WebSearchOptions, kwargs.get("web_search_options", {})) - return WebSearchOptions(**web_search_options) + web_search_options: Final = kwargs.get("web_search_options") + if isinstance(web_search_options, dict): + return WebSearchOptions(**cast(WebSearchOptions, web_search_options)) tools: Final = StandardBuiltInToolCostTracking._get_tools_from_kwargs( kwargs=kwargs, tool_type="web_search_preview" @@ -876,10 +876,9 @@ class StandardBuiltInToolCostTracking: _get_web_search_options = get_web_search_options @staticmethod - def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> list[object] | None: - if "tools" in kwargs: - return cast(list[object], kwargs.get("tools", [])) - return None + def _get_tools_from_kwargs(kwargs: Mapping[str, object], tool_type: str) -> Sequence[object] | None: + tools: Final = kwargs.get("tools") + return cast(Sequence[object], tools) if isinstance(tools, (list, tuple)) else None @staticmethod def get_file_search_tool_call(kwargs: Mapping[str, object]) -> FileSearchTool | None: diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py new file mode 100644 index 00000000000..f90c8f5d174 --- /dev/null +++ b/litellm/litellm_core_utils/tool_fields.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import reduce +from typing import Annotated, Final, TypeAlias + +from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.litellm_core_utils.core_helpers import normalize_drop_params + +ROUTES_WITH_TOOLS: Final = frozenset( + { + "acompletion", + "aresponses", + "acompact_responses", + "anthropic_messages", + "agenerate_content", + "agenerate_content_stream", + "acreate_interaction", + } +) +TOOL_FIELDS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls", "toolConfig", "tool_config"}) +_TOOLS: Final = TypeAdapter(list[object]) +_FIELD_LIST: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final = TypeAdapter(Json[dict[str, object]]) + + +def _schema_as_read(value: object) -> object: + if not value: + return None + if not isinstance(value, str): + return value + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return value + + +_Schema: TypeAlias = Annotated[dict[str, object] | None, BeforeValidator(_schema_as_read)] + + +class _ToolFunction(TypedDict, total=False): + parameters: ReadOnly[_Schema] + + +class _ToolEntry(TypedDict, total=False): + function: ReadOnly[_ToolFunction | None] + parameters: ReadOnly[_Schema] + input_schema: ReadOnly[_Schema] + vector_store_ids: ReadOnly[list[str] | None] + + +_TOOL_ENTRY: Final = TypeAdapter(_ToolEntry) +_WEB_SEARCH_OPTIONS: Final = TypeAdapter(dict[str, object] | None, config=ConfigDict(strict=True)) + + +@dataclass(frozen=True, slots=True) +class DropPolicy: + drop_params: bool + additional_drop_params: frozenset[str] + + def drops(self, field: str) -> bool: + return self.drop_params or field in self.additional_drop_params + + +@dataclass(frozen=True, slots=True) +class MalformedField: + field: str + message: str + dropped: Callable[[Mapping[str, object]], Mapping[str, object]] + + +def _field_names(value: object) -> frozenset[str]: + try: + return frozenset(name for name in _FIELD_LIST.validate_python(value) if isinstance(name, str)) + except ValidationError: + return frozenset() + + +def drop_policy(params: Mapping[str, object]) -> DropPolicy: + return DropPolicy( + drop_params=litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True, + additional_drop_params=_field_names(params.get("additional_drop_params")), + ) + + +def first_validation_error(e: ValidationError, prefix: tuple[str | int, ...] = ()) -> str: + error: Final = e.errors(include_url=False)[0] + return f"{'.'.join(str(part) for part in (*prefix, *error['loc']))}: {error['msg']}" + + +def _tool_entry_error(index: int, entry: object) -> str | None: + try: + _TOOL_ENTRY.validate_python(entry, strict=True) + except ValidationError as e: + return first_validation_error(e, prefix=("tools", index)) + return None + + +def _without_tool_fields(data: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in data.items() if name not in TOOL_FIELDS} + + +def _malformed_tools(data: Mapping[str, object]) -> MalformedField | None: + tools: Final = data.get("tools") + if tools is None: + return None + try: + entries: Final = _TOOLS.validate_python(tools) + except ValidationError as e: + return MalformedField("tools", first_validation_error(e, prefix=("tools",)), _without_tool_fields) + entry_errors: Final = tuple(_tool_entry_error(index, entry) for index, entry in enumerate(entries)) + first_error: Final = next((error for error in entry_errors if error is not None), None) + if first_error is None: + return None + kept: Final = [entry for entry, error in zip(entries, entry_errors) if error is None] + return MalformedField( + "tools", first_error, lambda fields: {**fields, "tools": kept} if kept else _without_tool_fields(fields) + ) + + +def _without_web_search_options(data: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in data.items() if name != "web_search_options"} + + +def _malformed_web_search_options(data: Mapping[str, object]) -> MalformedField | None: + try: + _WEB_SEARCH_OPTIONS.validate_python(data.get("web_search_options")) + except ValidationError as e: + return MalformedField( + "web_search_options", + first_validation_error(e, prefix=("web_search_options",)), + _without_web_search_options, + ) + return None + + +def resolve_tool_fields( + data: Mapping[str, object], candidates: Callable[[], tuple[DropPolicy, ...]] +) -> Mapping[str, object] | MalformedField: + checked: Final = (_malformed_tools(data), _malformed_web_search_options(data)) + malformed: Final = tuple(field for field in checked if field is not None) + if not malformed: + return data + policies: Final = candidates() + rejected: Final = next( + (field for field in malformed if not any(policy.drops(field.field) for policy in policies)), None + ) + if rejected is not None: + return rejected + return reduce( + lambda resolved, field: field.dropped(resolved), + (field for field in malformed if all(policy.drops(field.field) for policy in policies)), + data, + ) + + +def resolve_deployment_tool_fields( + kwargs: Mapping[str, object], route: object +) -> Mapping[str, object] | MalformedField: + if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS: + return kwargs + return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),)) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6bc1c4ead64..6b6974becac 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -45,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, @@ -924,10 +925,10 @@ class AnthropicMessagesHandler(BaseTranslation): images=(), ) - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract every tool name in an Anthropic messages request: tools[].name, plus tools[].function.name for OpenAI-format tools the bridge forwards verbatim.""" - return [name for tool in data.get("tools") or [] for name in anthropic_tool_names(tool)] + return [name for tool in list_field(data, "tools") for name in anthropic_tool_names(tool)] @classmethod def _extract_input_text_and_images( diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index 5fb2546a5ae..e65fe596da7 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -2,6 +2,7 @@ import asyncio import json import re from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence +from contextvars import ContextVar from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable @@ -352,6 +353,22 @@ async def aclose_if_supported(stream: object) -> None: await stream.aclose() +_CLOSING_SUPERSEDED_ATTEMPT: Final[ContextVar[bool]] = ContextVar("closing_superseded_attempt", default=False) + + +async def aclose_superseded_attempt(stream: object) -> None: + """Close an attempt a retry or fallback is taking over, so its end-of-stream logging records a failure.""" + token: Final = _CLOSING_SUPERSEDED_ATTEMPT.set(True) + try: + await aclose_if_supported(stream) + finally: + _CLOSING_SUPERSEDED_ATTEMPT.reset(token) + + +def closing_superseded_attempt() -> bool: + return _CLOSING_SUPERSEDED_ATTEMPT.get() + + _RESPONSE_HEADERS_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 18ecff990d2..692e5d4c175 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -43,6 +43,12 @@ _BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]]) _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() +def _record_bridged_response(responses_kwargs: Mapping[str, object], result: ResponsesAPIResponse) -> None: + logging_obj: Final = litellm_logging_obj_from_kwargs(responses_kwargs) + if logging_obj is not None: + logging_obj.record_bridged_responses_api_response(result) + + def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, object]: """The litellm-specific kwargs forwarded verbatim onto the Responses API request.""" return extra_kwargs or {} @@ -301,7 +307,9 @@ class LiteLLMMessagesToResponsesAPIHandler: if not isinstance(result, ResponsesAPIResponse): raise ValueError(f"Expected ResponsesAPIResponse, got {type(result)}") - return await with_safeguard_results(_ADAPTER.translate_response(result), safeguards_evaluator) + served: Final = await with_safeguard_results(_ADAPTER.translate_response(result), safeguards_evaluator) + _record_bridged_response(responses_kwargs, result) + return served @staticmethod def anthropic_messages_handler( diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 82150e415aa..a23fffb5744 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -95,7 +95,6 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): status_code=400, message=( "Azure OpenAI does not support reasoning_effort='none' for this model. " - "Supported values are: 'low', 'medium', and 'high'. " "To drop this parameter, set `litellm.drop_params=True` or for proxy:\n\n" "`litellm_settings:\n drop_params: true`\n" "Issue: https://github.com/BerriAI/litellm/issues/16704" diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 3564d7a1364..46106026920 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -2,6 +2,7 @@ import asyncio import hashlib import json import os +import re from collections.abc import Callable, Mapping from functools import lru_cache from types import MappingProxyType @@ -25,6 +26,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.utils import add_path_to_api_base azure_ad_cache: Final = DualCache() +_LAST_OPENAI_PATH_SEGMENT: Final = re.compile(r"^(.*/openai)(?=/|$)") class _AzureAdTokenJson(TypedDict, total=False): @@ -798,6 +800,10 @@ class BaseAzureLLM(BaseOpenAILLM): ) -> dict[str, str]: # mutable-ok: mirrors override contract return cls._base_validate_azure_environment(headers, litellm_params) + @staticmethod + def resolve_azure_api_base(api_base: str | None) -> str | None: + return api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + @staticmethod def _get_base_azure_url( api_base: str | None, @@ -815,7 +821,7 @@ class BaseAzureLLM(BaseOpenAILLM): default_api_version: The default API version to use if no api_version is provided. If 'latest', it will use `openai/v1/...` route. """ - api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + api_base = BaseAzureLLM.resolve_azure_api_base(api_base) if api_base is None: raise ValueError( f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`" @@ -843,7 +849,7 @@ class BaseAzureLLM(BaseOpenAILLM): # ensure the request go to /openai/v1 and not just /openai if "/openai/v1" not in new_url: parsed_url: Final = httpx.URL(new_url) - new_url = str(parsed_url.copy_with(path=parsed_url.path.replace("/openai", "/openai/v1"))) + new_url = str(parsed_url.copy_with(path=_LAST_OPENAI_PATH_SEGMENT.sub(r"\1/v1", parsed_url.path, 1))) # Use the new query_params dictionary final_url: Final = httpx.URL(new_url).copy_with(params=query_params) diff --git a/litellm/llms/azure/containers/transformation.py b/litellm/llms/azure/containers/transformation.py index 5031f38bf8f..d69c39d8873 100644 --- a/litellm/llms/azure/containers/transformation.py +++ b/litellm/llms/azure/containers/transformation.py @@ -28,6 +28,7 @@ class AzureContainerConfig(OpenAIContainerConfig): self, headers: dict, api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, ) -> dict: return BaseAzureLLM.base_validate_azure_environment( headers=headers, diff --git a/litellm/llms/azure/fine_tuning/handler.py b/litellm/llms/azure/fine_tuning/handler.py index 0baf3ccc598..aa410d8232b 100644 --- a/litellm/llms/azure/fine_tuning/handler.py +++ b/litellm/llms/azure/fine_tuning/handler.py @@ -1,4 +1,4 @@ -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from typing import Final, cast import httpx @@ -66,6 +66,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: self._ensure_training_type(create_fine_tuning_job_data) @@ -78,6 +79,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -109,6 +111,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -119,6 +122,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -149,6 +153,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -159,6 +164,7 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -188,14 +194,14 @@ class AzureOpenAIFineTuningAPI(OpenAIFineTuningAPI, BaseAzureLLM): client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, _is_async: bool = False, api_version: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None: # Override to use Azure-specific client initialization if isinstance(client, OpenAI) or isinstance(client, AsyncOpenAI): client = None return self.get_azure_openai_client( - litellm_params=litellm_params or {}, + litellm_params=dict(litellm_params) if litellm_params is not None else None, api_key=api_key, api_base=api_base, api_version=api_version, diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 98b4241d7ab..75714d5253b 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -1,4 +1,7 @@ +import re +from collections.abc import Mapping from copy import deepcopy +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal import httpx @@ -6,6 +9,7 @@ from openai.types.responses import ResponseReasoningItem from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM @@ -25,6 +29,34 @@ else: _DECODED_JSON: Final = TypeAdapter(object) +_DATED_GA_API_VERSION: Final = re.compile(r"\d{4}-\d{2}-\d{2}") + + +def _str_param(litellm_params: Mapping[str, object], key: str) -> str | None: + value: Final = litellm_params.get(key) + return value if isinstance(value, str) and value else None + + +def _resolve_responses_api_version(api_base: str | None, litellm_params: Mapping[str, object]) -> str | None: + """ + Dated GA api_versions (e.g. 2024-10-21) predate the Responses API on Azure and 404 there, so a deployment + pinned to one for chat gets the default Responses version unless it sets responses_api_version. + """ + api_version: Final = _str_param(litellm_params, "api_version") + if api_base is not None and "api-version" in httpx.URL(api_base).params: + return api_version + responses_api_version: Final = _str_param(litellm_params, "responses_api_version") + if responses_api_version is not None: + return responses_api_version + if api_version is not None and _DATED_GA_API_VERSION.fullmatch(api_version): + return AZURE_DEFAULT_RESPONSES_API_VERSION + return api_version + + +def _with_responses_api_version(api_base: str | None, litellm_params: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({**litellm_params, "api_version": _resolve_responses_api_version(api_base, litellm_params)}) + + class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Parameters not supported by Azure Responses API AZURE_UNSUPPORTED_PARAMS = ["context_management"] @@ -170,11 +202,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - A complete URL string, e.g., "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2024-05-01-preview" """ - from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION - + resolved_api_base: Final = BaseAzureLLM.resolve_azure_api_base(api_base) return BaseAzureLLM.get_base_azure_url( - api_base=api_base, - litellm_params=litellm_params, + api_base=resolved_api_base, + litellm_params=_with_responses_api_version(resolved_api_base, litellm_params), route="/openai/responses", default_api_version=AZURE_DEFAULT_RESPONSES_API_VERSION, ) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index a8e6320bd1e..fc22e9a8b73 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,6 +1,7 @@ import copy import enum import re +from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Final, cast import httpx @@ -21,12 +22,12 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error -from litellm.llms.openai.openai import OpenAIConfig +from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ModelResponse, ProviderField +from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField from litellm.utils import add_path_to_api_base, supports_tool_choice if TYPE_CHECKING: @@ -45,6 +46,18 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) +def _is_grok_model(model: str) -> bool: + return "grok" in model + + +class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): + def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract + usage: Final = chunk.get("usage") + if _is_grok_model(str(chunk.get("model") or "")) and isinstance(usage, dict): + XAIChatConfig.fold_reasoning_tokens_into_completion(usage) + return super().chunk_parser(chunk) + + class AzureAIGPT5Config(OpenAIGPT5Config): @classmethod def _model_map_lookup_name(cls, model: str) -> str: @@ -114,7 +127,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Check if the model supports stop tokens. """ - if "grok" in model: + if _is_grok_model(model): # Reuse Xai method for Grok model xai_config: Final = XAIChatConfig() return xai_config.supports_stop_reason(model) @@ -319,7 +332,7 @@ class AzureAIStudioConfig(OpenAIConfig): json_mode: bool | None = None, ) -> ModelResponse: model_response.model = f"azure_ai/{model}" - return super().transform_response( + response: Final = super().transform_response( model=model, raw_response=raw_response, model_response=model_response, @@ -332,6 +345,21 @@ class AzureAIStudioConfig(OpenAIConfig): api_key=api_key, json_mode=json_mode, ) + if _is_grok_model(model): + XAIChatConfig.fold_reasoning_tokens_into_completion(response) + return response + + def get_model_response_iterator( + self, + streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, + sync_stream: bool, + json_mode: bool | None = False, + ) -> "AzureAIChatCompletionResponseIterator": + return AzureAIChatCompletionResponseIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) def should_retry_llm_api_inside_llm_translation_on_http_error( self, e: httpx.HTTPStatusError, litellm_params: dict diff --git a/litellm/llms/base_llm/containers/transformation.py b/litellm/llms/base_llm/containers/transformation.py index 49487cddb6f..7a04c4d9488 100644 --- a/litellm/llms/base_llm/containers/transformation.py +++ b/litellm/llms/base_llm/containers/transformation.py @@ -81,6 +81,7 @@ class BaseContainerConfig(ABC): self, headers: dict, api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, ) -> dict: return {} diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 041473cb634..df7b7645840 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -495,6 +495,11 @@ class BaseDecisionsConfig(ABC): def resolve_connection(self, api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]: return self.resolve_api_base(api_base), self.resolve_api_key(api_key) + def resolve_credential( + self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object] + ) -> str | None: + return api_key + @staticmethod def _first_secret(names: tuple[str, ...]) -> str | None: return next((value for value in (get_secret_str(name) for name in names) if value), None) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 39ff2fa0528..55a4b3da0b3 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,7 +4,7 @@ import json from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -292,6 +292,17 @@ def openai_tool_name(tool: object) -> str | None: return flat_name if isinstance(flat_name, str) else None +_FIELDS: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(strict=True)) +_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) + + +def list_field(data: object, key: str) -> Sequence[object]: + try: + return _ITEMS.validate_python(_FIELDS.validate_python(data).get(key)) + except ValidationError: + return () + + def anthropic_tool_names(tool: object) -> tuple[str, ...]: """Every name a /v1/messages tool dict can act under: the flat Anthropic ``name`` plus ``function.name`` for OpenAI-format tools the bridge forwards verbatim. Allowlist checks diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 04747405375..fc6088ddc66 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -32,6 +32,7 @@ from litellm.llms.bedrock.common_utils import ( tools_without_eager_input_streaming, ) from litellm.types.llms.anthropic import ( + ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, AnthropicThinkingParam, @@ -270,7 +271,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages) is_thinking_display_updates_used: Final = self.is_thinking_display_updates_used(thinking) - user_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) + client_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) + user_beta_set: Final = ( + client_beta_set + if "safeguards" in optional_params + else client_beta_set - {ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value} + ) beta_set: Final = set(user_beta_set) auto_betas: Final = self.get_anthropic_beta_list( model=model, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 88ae0d4d514..95e81fa3ee7 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -90,6 +90,9 @@ class AmazonAnthropicClaudeMessagesConfig( def beta_headers_provider(self) -> str: return self.custom_llm_provider or "bedrock" + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: + return True + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) def get_error_class( @@ -552,8 +555,17 @@ class AmazonAnthropicClaudeMessagesConfig( if anthropic_model_info.is_eager_input_streaming_used(tools): beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + dangerous_tool_use_beta: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value if anthropic_messages_optional_request_params.get("safeguards") is not None: - beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + beta_set.add(dangerous_tool_use_beta) + elif dangerous_tool_use_beta in beta_set and self.drops_dangerous_tool_use_beta_without_safeguards(): + beta_set.discard(dangerous_tool_use_beta) + verbose_logger.debug( + "Bedrock Invoke: dropping %s from client headers because the request has no safeguards. Bedrock " + "answers it with 'invalid beta flag' next to structured outputs before Opus 4.7, and on Sonnet 4.5, " + "Opus 4.5, and Opus 4.6 even alone.", + dangerous_tool_use_beta, + ) self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index ca93cc3e617..1ad84761bc9 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -67,6 +67,9 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: + return False + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index bfa163c7546..a9b96217379 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -5,6 +5,7 @@ This module provides a single generic handler that can process any container fil endpoint defined in endpoints.json, eliminating the need for individual handler methods. """ +import asyncio import json from collections.abc import Coroutine, Mapping, Sequence from pathlib import Path @@ -236,6 +237,7 @@ def _request_headers( return container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) @@ -435,7 +437,7 @@ class GenericContainerHandler: http_client: Final = _async_http_client(client, litellm_params) # Build request - headers = _request_headers(container_provider_config, extra_headers, litellm_params) + headers = await asyncio.to_thread(_request_headers, container_provider_config, extra_headers, litellm_params) if extra_headers: headers.update(extra_headers) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 70b4d2d0f36..8039b85d964 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7129,7 +7129,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = image_edit_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + image_edit_provider_config.validate_environment, api_key=litellm_params.api_key, headers=image_edit_optional_request_params.get("extra_headers", {}) or {}, model=model, @@ -7588,7 +7589,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_generation_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_generation_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=video_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, @@ -7783,7 +7785,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_content_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_content_provider_config.validate_environment, headers=extra_headers or {}, model="", api_key=api_key, @@ -7963,7 +7966,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_remix_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_remix_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", @@ -8126,7 +8130,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", @@ -8266,7 +8271,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", @@ -8438,7 +8444,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", @@ -8624,7 +8631,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_provider_config.validate_environment, api_key=api_key or litellm_params.get("api_key", None), headers=extra_headers or {}, model="", @@ -8756,7 +8764,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_list_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_list_provider_config.validate_environment, api_key=api_key, headers=extra_headers or {}, model="", @@ -8837,7 +8846,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_delete_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_delete_provider_config.validate_environment, api_key=api_key, headers=extra_headers or {}, model="", @@ -9021,7 +9031,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = video_status_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + video_status_provider_config.validate_environment, api_key=api_key, headers=extra_headers or {}, model="", @@ -9124,6 +9135,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) # Add Content-Type header for JSON requests @@ -9199,9 +9211,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) # Add Content-Type header for JSON requests @@ -9295,6 +9309,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9371,9 +9386,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9462,6 +9479,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9538,9 +9556,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9631,6 +9651,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9707,9 +9728,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9805,6 +9828,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9883,9 +9907,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -9974,6 +10000,7 @@ class BaseLLMHTTPHandler: headers: Final = container_provider_config.validate_environment( headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -10049,9 +10076,11 @@ class BaseLLMHTTPHandler: async_httpx_client = client # Validate environment and get headers - headers: Final = container_provider_config.validate_environment( + headers: Final = await asyncio.to_thread( + container_provider_config.validate_environment, headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), + litellm_params=litellm_params, ) if extra_headers: @@ -10169,8 +10198,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -10412,8 +10443,10 @@ class BaseLLMHTTPHandler: vector_store_provider_config.validate_create_vector_store() - headers: Final = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -10543,8 +10576,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -10665,8 +10700,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -10807,8 +10844,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -10952,8 +10991,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -11070,8 +11111,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: @@ -11212,8 +11255,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -11353,8 +11398,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -11484,8 +11531,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -11621,8 +11670,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) @@ -11764,8 +11815,10 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = vector_store_files_provider_config.validate_environment( - headers=extra_headers or {}, litellm_params=litellm_params + headers: Final = await asyncio.to_thread( + vector_store_files_provider_config.validate_environment, + headers=extra_headers or {}, + litellm_params=litellm_params, ) if extra_headers: headers.update(extra_headers) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index c168c3f88fc..3070bedd30d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, openai_tool_name, @@ -156,7 +157,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices] # Pass tools (function definitions) to the guardrail tools: Final = data.get("tools") - if tools and not scan_only_tool_results: + if isinstance(tools, list) and tools and not scan_only_tool_results: inputs["tools"] = tools # Include model information if available model: Final = data.get("model") @@ -258,19 +259,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return "no scannable content after message scoping" return None if images else "no scannable content" - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name, and the tools a replayed tool_search_output loads).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in list_field(data, "tools"): if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or []: + for fn in list_field(data, "functions"): if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - return [*names, *replayed_tool_search_tool_names(cast(object, data))] # cast-ok: upcast, read field by field + return [*names, *replayed_tool_search_tool_names(data)] def _extract_inputs( self, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 4ab91392ad4..e472661efb9 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -407,6 +407,10 @@ class BaseOpenAILLM: http2=http2_enabled(), ) + @classmethod + def get_sync_http_client(cls) -> httpx.Client | None: + return cls._get_sync_http_client() + class OpenAICredentials(NamedTuple): api_base: str diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 174aefd29e5..278427e8d77 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -1,9 +1,7 @@ import json -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Final -from openai import AsyncOpenAI, OpenAI - import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -12,7 +10,8 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse from litellm.utils import ProviderConfigManager -from ..common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient +from ..common_utils import BaseOpenAILLM, OpenAIError +from ..workload_identity import build_async_openai_client, build_openai_client from .transformation import OpenAITextCompletionConfig @@ -43,7 +42,7 @@ class OpenAITextCompletion(BaseLLM): print_verbose: Callable | None = None, api_base: str | None = None, acompletion: bool = False, - litellm_params=None, + litellm_params: Mapping[str, object] | None = None, logger_fn=None, client=None, organization: str | None = None, @@ -95,6 +94,7 @@ class OpenAITextCompletion(BaseLLM): max_retries=max_retries, client=client, organization=organization, + litellm_params=litellm_params, ) else: return self.acompletion( @@ -109,6 +109,7 @@ class OpenAITextCompletion(BaseLLM): max_retries=max_retries, organization=organization, client=client, + litellm_params=litellm_params, ) elif optional_params.get("stream", False): return self.streaming( @@ -123,19 +124,22 @@ class OpenAITextCompletion(BaseLLM): max_retries=max_retries, client=client, organization=organization, + litellm_params=litellm_params, ) else: - if client is None: - openai_client = OpenAI( + openai_client: Final = ( + client + if client is not None + else build_openai_client( api_key=api_key, - base_url=api_base, - http_client=litellm.client_session or OpenAIHTTPClient(), + api_base=api_base, timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=lambda: litellm.client_session, ) - else: - openai_client = client + ) raw_response: Final = openai_client.completions.with_raw_response.create(**data) response: Final = raw_response.parse() @@ -175,19 +179,22 @@ class OpenAITextCompletion(BaseLLM): max_retries: int, organization: str | None = None, client=None, + litellm_params: Mapping[str, object] | None = None, ): try: - if client is None: - openai_aclient = AsyncOpenAI( + openai_aclient: Final = ( + client + if client is not None + else build_async_openai_client( api_key=api_key, - base_url=api_base, - http_client=BaseOpenAILLM.get_async_http_client(), + api_base=api_base, timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=BaseOpenAILLM.get_async_http_client, ) - else: - openai_aclient = client + ) raw_response: Final = await openai_aclient.completions.with_raw_response.create(**data) response: Final = raw_response.parse() @@ -228,18 +235,21 @@ class OpenAITextCompletion(BaseLLM): max_retries=None, client=None, organization=None, + litellm_params: Mapping[str, object] | None = None, ): - if client is None: - openai_client = OpenAI( + openai_client: Final = ( + client + if client is not None + else build_openai_client( api_key=api_key, - base_url=api_base, - http_client=litellm.client_session or OpenAIHTTPClient(), + api_base=api_base, timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=lambda: litellm.client_session, ) - else: - openai_client = client + ) try: raw_response: Final = openai_client.completions.with_raw_response.create(**data) @@ -285,18 +295,21 @@ class OpenAITextCompletion(BaseLLM): api_base: str | None = None, client=None, organization=None, + litellm_params: Mapping[str, object] | None = None, ): - if client is None: - openai_client = AsyncOpenAI( + openai_client: Final = ( + client + if client is not None + else build_async_openai_client( api_key=api_key, - base_url=api_base, - http_client=litellm.aclient_session or OpenAIAsyncHTTPClient(), + api_base=api_base, timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=lambda: litellm.aclient_session, ) - else: - openai_client = client + ) raw_response: Final = await openai_client.completions.with_raw_response.create(**data) response: Final = raw_response.parse() diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index bff5e0cea10..6cb2396b85d 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -23,6 +23,7 @@ from litellm.types.containers.main import ( from litellm.types.router import GenericLiteLLMParams from ...base_llm.containers.transformation import BaseContainerConfig +from ..workload_identity import resolve_openai_bearer_token from .utils import join_container_api_base_path if TYPE_CHECKING: @@ -104,13 +105,15 @@ class OpenAIContainerConfig(BaseContainerConfig): self, headers: dict, api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, ) -> dict: - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - headers.update( - { - "Authorization": f"Bearer {api_key}", - }, + params: Final = litellm_params or GenericLiteLLMParams() + bearer_token: Final = resolve_openai_bearer_token( + api_key=api_key or params.api_key, + api_base=params.api_base, + litellm_params=params.model_dump(exclude_none=True), ) + headers.update({"Authorization": f"Bearer {bearer_token}"}) return headers def get_complete_url( diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 7ce5b8676ee..56f348881de 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -6,6 +6,7 @@ from typing_extensions import assert_never import litellm from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, decisions_text +from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.secret_managers.main import get_secret_str from litellm.types.decisions import ( DecisionsImage, @@ -349,6 +350,11 @@ class OpenAIDecisionsConfig(BaseDecisionsConfig): def resolve_api_key(self, api_key: str | None) -> str | None: return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + def resolve_credential( + self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object] + ) -> str | None: + return resolve_openai_bearer_token(api_key=api_key, api_base=api_base, litellm_params=litellm_params) + def supports_audio_input(self, model: str, custom_llm_provider: str) -> bool: return False diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 7980c5529cb..24de7f6d479 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -13,6 +13,7 @@ from litellm.llms.base_llm.evals.transformation import ( BaseEvalsAPIConfig, LiteLLMLoggingObj, ) +from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.types.llms.openai_evals import ( CancelEvalResponse, CancelRunResponse, @@ -44,24 +45,14 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): return LlmProviders.OPENAI def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: - """Add OpenAI-specific headers""" - import litellm - from litellm.secret_managers.main import get_secret_str - - # Get API key following OpenAI pattern - api_key = None - if litellm_params: - api_key = litellm_params.api_key - - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - - if not api_key: + params: Final = litellm_params or GenericLiteLLMParams() + bearer_token: Final = resolve_openai_bearer_token( + api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True) + ) + if not bearer_token: raise ValueError("OPENAI_API_KEY is required for Evals API") - - # Add required headers - headers["Authorization"] = f"Bearer {api_key}" + headers["Authorization"] = f"Bearer {bearer_token}" headers["Content-Type"] = "application/json" - return headers def get_complete_url( diff --git a/litellm/llms/openai/fine_tuning/handler.py b/litellm/llms/openai/fine_tuning/handler.py index 851113bbd37..9ba28803f93 100644 --- a/litellm/llms/openai/fine_tuning/handler.py +++ b/litellm/llms/openai/fine_tuning/handler.py @@ -6,7 +6,7 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from openai.types.fine_tuning import FineTuningJob from litellm._logging import verbose_logger -from litellm.llms.openai.common_utils import OpenAIAsyncHTTPClient, OpenAIHTTPClient +from litellm.llms.openai.workload_identity import build_async_openai_client, build_openai_client from litellm.types.utils import LiteLLMFineTuningJob _AZURE_STATUS_MAP: Final[Mapping[object, str]] = { @@ -74,27 +74,27 @@ class OpenAIFineTuningAPI: client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, _is_async: bool = False, api_version: str | None = None, - litellm_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None: - received_args: Final = locals() - openai_client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client" or k == "_is_async": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - if _is_async is True: - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + if _is_async: + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) async def acreate_fine_tuning_job( self, @@ -116,6 +116,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -126,6 +127,7 @@ class OpenAIFineTuningAPI: client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -164,6 +166,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -174,6 +177,7 @@ class OpenAIFineTuningAPI: client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -212,6 +216,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, after: str | None = None, limit: int | None = None, ): @@ -224,6 +229,7 @@ class OpenAIFineTuningAPI: client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -263,6 +269,7 @@ class OpenAIFineTuningAPI: max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]: openai_client: Final[OpenAI | AsyncOpenAI | AzureOpenAI | AsyncAzureOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -273,6 +280,7 @@ class OpenAIFineTuningAPI: client=client, _is_async=_is_async, api_version=api_version, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index d54522597a0..a9e5b1b9f34 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -17,6 +17,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.utils import ImageResponse from ..common_utils import OpenAIError +from ..workload_identity import resolve_openai_bearer_token if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -160,12 +161,10 @@ class OpenAIImageEditConfig(BaseImageEditConfig): litellm_params: dict | None = None, api_base: str | None = None, ) -> dict: - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - headers.update( - { - "Authorization": f"Bearer {api_key}", - } + bearer_token: Final = resolve_openai_bearer_token( + api_key=api_key, api_base=api_base, litellm_params=litellm_params ) + headers.update({"Authorization": f"Bearer {bearer_token}"}) return headers def get_complete_url( diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 664d16c48bf..89dcdc40e02 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -2,7 +2,7 @@ OpenAI Image Variations Handler """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Final import httpx @@ -14,27 +14,58 @@ from litellm.utils import ProviderConfigManager from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj -from ..common_utils import OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient +from ..common_utils import OpenAIError +from ..workload_identity import build_async_openai_client, build_openai_client class OpenAIImageVariationsHandler: def get_sync_client( self, client: OpenAI | None, - init_client_params: dict, - ): + api_key: str | None, + api_base: str, + timeout: float | httpx.Timeout | None, + max_retries: int, + organization: str | None, + litellm_params: Mapping[str, object], + ) -> OpenAI: if client is not None: return client - return OpenAI(**init_client_params, http_client=litellm.client_session or OpenAIHTTPClient()) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=lambda: litellm.client_session, + ) - def get_async_client(self, client: AsyncOpenAI | None, init_client_params: dict) -> AsyncOpenAI: + def get_async_client( + self, + client: AsyncOpenAI | None, + api_key: str | None, + api_base: str, + timeout: float | httpx.Timeout | None, + max_retries: int, + organization: str | None, + litellm_params: Mapping[str, object], + ) -> AsyncOpenAI: if client is not None: return client - return AsyncOpenAI(**init_client_params, http_client=litellm.aclient_session or OpenAIAsyncHTTPClient()) + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + static_key_http_client_factory=lambda: litellm.aclient_session, + ) async def async_image_variations( self, - api_key: str, + api_key: str | None, api_base: str, organization: str | None, client: AsyncOpenAI | None, @@ -51,15 +82,15 @@ class OpenAIImageVariationsHandler: provider_config: BaseImageVariationConfig, ) -> ImageResponse: try: - init_client_params: Final = { - "api_key": api_key, - "base_url": api_base, - "timeout": timeout, - "max_retries": max_retries, - "organization": organization, - } - - client = self.get_async_client(client=client, init_client_params=init_client_params) + client = self.get_async_client( + client=client, + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) raw_response: Final = await client.images.with_raw_response.create_variation(**data) response: Final = raw_response.parse() @@ -103,7 +134,7 @@ class OpenAIImageVariationsHandler: def image_variations( self, model_response: ImageResponse, - api_key: str, + api_key: str | None, api_base: str, model: str | None, image: FileTypes, @@ -167,15 +198,15 @@ class OpenAIImageVariationsHandler: litellm_params=litellm_params, ) - init_client_params: Final = { - "api_key": api_key, - "base_url": api_base, - "timeout": timeout, - "max_retries": max_retries, - "organization": organization, - } - - client = self.get_sync_client(client=client, init_client_params=init_client_params) + client = self.get_sync_client( + client=client, + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) raw_response: Final = client.images.with_raw_response.create_variation(**json_data) response: Final = raw_response.parse() diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 79c69b5ceb4..0c1eaaaae26 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -52,16 +52,19 @@ from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, - OpenAIAsyncHTTPClient, OpenAIError, - OpenAIHTTPClient, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, is_openai_backed_api_base, is_output_token_limit_error, model_dump_provider_response, ) -from .workload_identity import resolve_openai_workload_identity_config +from .workload_identity import ( + build_async_openai_client, + build_openai_client, + openai_static_api_key, + resolve_openai_workload_identity_config, +) openaiOSeriesConfig: Final = OpenAIOSeriesConfig() openAIGPT5Config: Final = OpenAIGPT5Config() @@ -425,7 +428,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else AsyncOpenAI( - api_key=api_key, + api_key=openai_static_api_key(api_key), base_url=api_base, http_client=async_http_client, timeout=timeout, @@ -447,7 +450,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else OpenAI( - api_key=api_key, + api_key=openai_static_api_key(api_key), base_url=api_base, http_client=sync_http_client, timeout=timeout, @@ -1455,10 +1458,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=None, organization: str | None = None, headers: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ): try: openai_aclient: Final = self._get_openai_client( is_async=True, + litellm_params=litellm_params, api_key=api_key, api_base=api_base, timeout=timeout, @@ -1520,6 +1525,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aimg_generation=None, organization: str | None = None, headers: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> ImageResponse: data = {} try: @@ -1541,10 +1547,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, organization=organization, headers=headers, + litellm_params=litellm_params, ) openai_client: Final[OpenAI] = self._get_openai_client( is_async=False, + litellm_params=litellm_params, api_key=api_key, api_base=api_base, timeout=timeout, @@ -1622,6 +1630,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aspeech: bool | None = None, client=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent: if aspeech is not None and aspeech is True: return self.async_audio_speech( @@ -1638,6 +1647,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj=logging_obj, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) openai_client: Final = self._get_openai_client( @@ -1648,6 +1658,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) sync_client: Final = cast(OpenAI, openai_client) @@ -1685,6 +1696,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, client=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent: openai_client: Final = cast( AsyncOpenAI, @@ -1696,6 +1708,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ), ) @@ -1742,26 +1755,27 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - openai_client: OpenAI | AsyncOpenAI | None = None - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client" or k == "_is_async": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - if _is_async is True: - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + if _is_async: + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) async def acreate_file( self, @@ -1781,6 +1795,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1790,6 +1805,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1823,6 +1839,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1832,6 +1849,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1888,6 +1906,7 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, chunk_size: int = 1024 * 1024, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> FileContentStreamingResult: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1897,6 +1916,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1951,6 +1971,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1960,6 +1981,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1997,6 +2019,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2006,6 +2029,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2046,6 +2070,7 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, purpose: str | None = None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2055,6 +2080,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2100,26 +2126,27 @@ class OpenAIBatchesAPI(BaseLLM): organization: str | None, client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - openai_client: OpenAI | AsyncOpenAI | None = None - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client" or k == "_is_async": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - if _is_async is True: - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + if _is_async: + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) async def acreate_batch( self, @@ -2139,6 +2166,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2148,6 +2176,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2183,6 +2212,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2192,6 +2222,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2226,6 +2257,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2235,6 +2267,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2275,6 +2308,7 @@ class OpenAIBatchesAPI(BaseLLM): after: str | None = None, limit: int | None = None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2284,6 +2318,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2312,22 +2347,19 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) def async_get_openai_client( self, @@ -2337,22 +2369,19 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncOpenAI: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) ### ASSISTANTS ### @@ -2368,6 +2397,8 @@ class OpenAIAssistantsAPI(BaseLLM): limit: int | None = 20, before: str | None = None, after: str | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncCursorPage[Assistant]: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2375,6 +2406,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) request_params: Final = { @@ -2402,6 +2434,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_assistants: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, AsyncCursorPage[Assistant]]: ... @@ -2415,6 +2449,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_assistants: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[Assistant]: ... @@ -2433,7 +2469,9 @@ class OpenAIAssistantsAPI(BaseLLM): limit: int | None = 20, before: str | None = None, after: str | None = None, - ): + *, + litellm_params: Mapping[str, object] | None = None, + ) -> Coroutine[None, None, AsyncCursorPage[Assistant]] | SyncCursorPage[Assistant]: if aget_assistants is not None and aget_assistants is True: return self.async_get_assistants( api_key=api_key, @@ -2441,6 +2479,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2449,6 +2488,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2476,6 +2516,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, create_assistant_data: dict, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Assistant: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2483,6 +2525,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2500,6 +2543,8 @@ class OpenAIAssistantsAPI(BaseLLM): create_assistant_data: dict, client=None, async_create_assistants=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if async_create_assistants is not None and async_create_assistants is True: return self.async_create_assistants( @@ -2508,6 +2553,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, create_assistant_data=create_assistant_data, ) @@ -2517,6 +2563,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2533,6 +2580,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, assistant_id: str, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AssistantDeleted: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2540,6 +2589,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2557,6 +2607,8 @@ class OpenAIAssistantsAPI(BaseLLM): assistant_id: str, client=None, async_delete_assistants=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if async_delete_assistants is not None and async_delete_assistants is True: return self.async_delete_assistant( @@ -2565,6 +2617,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, assistant_id=assistant_id, ) @@ -2574,6 +2627,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2592,6 +2646,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIMessage: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2599,6 +2655,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2629,6 +2686,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, a_add_message: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, OpenAIMessage]: ... @@ -2644,6 +2703,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, a_add_message: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIMessage: ... @@ -2660,6 +2721,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, a_add_message: bool | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ): if a_add_message is not None and a_add_message is True: return self.a_add_message( @@ -2670,6 +2733,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2678,6 +2742,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2703,6 +2768,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncCursorPage[OpenAIMessage]: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2710,6 +2777,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2730,6 +2798,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_messages: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, AsyncCursorPage[OpenAIMessage]]: ... @@ -2744,6 +2814,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_messages: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[OpenAIMessage]: ... @@ -2759,6 +2831,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, aget_messages=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if aget_messages is not None and aget_messages is True: return self.async_get_messages( @@ -2768,6 +2842,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2776,6 +2851,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2795,6 +2871,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, messages: Iterable[OpenAICreateThreadParamsMessage] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2802,6 +2880,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2834,6 +2913,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client: AsyncOpenAI | None, acreate_thread: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Thread]: ... @@ -2849,6 +2930,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client: OpenAI | None, acreate_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2865,6 +2948,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client=None, acreate_thread=None, + *, + litellm_params: Mapping[str, object] | None = None, ): """ Here's an example: @@ -2884,6 +2969,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, messages=messages, ) @@ -2893,6 +2979,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2920,6 +3007,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2927,6 +3016,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2952,6 +3042,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_thread: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Thread]: ... @@ -2966,6 +3058,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2981,6 +3075,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, aget_thread=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if aget_thread is not None and aget_thread is True: return self.async_get_thread( @@ -2990,6 +3086,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2998,6 +3095,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3031,6 +3129,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Run: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -3038,6 +3138,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3142,6 +3243,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[True], event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Run]: ... @@ -3164,6 +3267,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[False] | None, event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Run: ... @@ -3187,6 +3292,8 @@ class OpenAIAssistantsAPI(BaseLLM): client=None, arun_thread=None, event_handler: AssistantEventHandler | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ): if arun_thread is not None and arun_thread is True: if stream is not None and stream is True: @@ -3196,6 +3303,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) return self.async_run_thread_stream( @@ -3223,6 +3331,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -3231,6 +3340,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 274d693baa7..c7c997bc961 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -27,6 +27,7 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( from litellm.llms.openai.workload_identity import ( OpenAIWorkloadIdentityConfig, get_workload_identity_bearer_token, + openai_static_api_key, resolve_openai_workload_identity_config, ) from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse @@ -35,13 +36,18 @@ from litellm.types.utils import CountContextManagement, LlmProviders, TokenCount openai_count_tokens_handler: Final = OpenAICountTokensHandler() -def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: +def _workload_identity_config( + litellm_params: Mapping[str, object], +) -> OpenAIWorkloadIdentityConfig | OpenAIError | None: api_base: Final = litellm_params.get("api_base") - return resolve_openai_workload_identity_config( - api_key=None, - api_base=api_base if isinstance(api_base, str) else None, - litellm_params=litellm_params, - ) + try: + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + litellm_params=litellm_params, + ) + except OpenAIError as e: + return e class OpenAITokenCounter(BaseTokenCounter): @@ -126,7 +132,18 @@ class OpenAITokenCounter(BaseTokenCounter): credentials: Final = get_openai_credentials( api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") ) - credential: Final = credentials.api_key or _workload_identity_config(litellm_params) + credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params) + + if isinstance(credential, OpenAIError): + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=credential.message, + status_code=credential.status_code, + ) if credential is None: return TokenCountResponse( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 435c11c0602..84ff790ae48 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -55,6 +55,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, effective_skip_system_message_for_guardrail, + list_field, merge_guardrailed_scoped_messages, role_out_of_guardrail_scope, scoped_structured_message_indices, @@ -764,11 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation): data.get("input"), data.get("instructions"), structured_messages, merged ) - def extract_request_tool_names(self, data: dict) -> list[str]: + def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: """Extract tool names from Responses API request (tools[].name for function and custom, tools[].server_label for mcp).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in list_field(data, "tools"): if not isinstance(tool, dict): continue if tool.get("type") in ("function", "custom") and tool.get("name"): diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 733d7930144..48a44575989 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast import httpx @@ -97,6 +98,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): max_retries=max_retries, logging_obj=logging_obj, shared_session=shared_session, + litellm_params=litellm_params, ) openai_client: Final[OpenAI] = self._get_openai_client( @@ -106,6 +108,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): timeout=timeout, max_retries=max_retries, client=client, + litellm_params=litellm_params, ) ## LOGGING @@ -159,6 +162,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): client=None, max_retries=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ): try: openai_aclient: Final[AsyncOpenAI] = self._get_openai_client( @@ -169,6 +173,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) ## LOGGING diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index d5f15d639d6..7f65ade224e 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) +from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_store_files import ( @@ -70,11 +71,13 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): headers: dict[str, str], litellm_params: GenericLiteLLMParams | None, ) -> dict[str, str]: - litellm_params = litellm_params or GenericLiteLLMParams() - api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + params: Final = litellm_params or GenericLiteLLMParams() + bearer_token: Final = resolve_openai_bearer_token( + api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True) + ) headers.update( { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json", } ) diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 57fe5b04838..46105a66be3 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -5,6 +5,7 @@ import httpx import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -48,11 +49,13 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): } def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict: - litellm_params = litellm_params or GenericLiteLLMParams() - api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + params: Final = litellm_params or GenericLiteLLMParams() + bearer_token: Final = resolve_openai_bearer_token( + api_key=params.api_key, api_base=params.api_base, litellm_params=params.model_dump(exclude_none=True) + ) headers.update( { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json", } ) diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 6e28695a96f..d422e51f649 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -8,11 +8,10 @@ from urllib.parse import quote import httpx from httpx._types import FileContent, FileTypes, RequestFiles -import litellm from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils -from litellm.secret_managers.main import get_secret_str +from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.types.llms.openai import CreateVideoRequest from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import ( @@ -78,16 +77,13 @@ class OpenAIVideoConfig(BaseVideoConfig): api_key: str | None = None, litellm_params: GenericLiteLLMParams | None = None, ) -> dict: - # Use api_key from litellm_params if available, otherwise fall back to other sources - if litellm_params and litellm_params.api_key: - api_key = api_key or litellm_params.api_key - - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - headers.update( - { - "Authorization": f"Bearer {api_key}", - } + params: Final = litellm_params or GenericLiteLLMParams() + bearer_token: Final = resolve_openai_bearer_token( + api_key=api_key or params.api_key, + api_base=params.api_base, + litellm_params=params.model_dump(exclude_none=True), ) + headers.update({"Authorization": f"Bearer {bearer_token}"}) return headers def get_complete_url( diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index e2971a2c100..5952ae3b3ef 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -6,10 +6,22 @@ from functools import lru_cache from typing import TYPE_CHECKING, Final from urllib.parse import urlparse -import litellm -from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str +import httpx +from openai import DEFAULT_MAX_RETRIES as OPENAI_SDK_DEFAULT_MAX_RETRIES +from openai import NOT_GIVEN, AsyncOpenAI, NotGiven, OpenAI -from .common_utils import OpenAIError, is_openai_backed_api_base +import litellm +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE +from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str +from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS + +from .common_utils import ( + BaseOpenAILLM, + OpenAIAsyncHTTPClient, + OpenAIError, + OpenAIHTTPClient, + is_openai_backed_api_base, +) if TYPE_CHECKING: from collections.abc import Callable @@ -17,6 +29,11 @@ if TYPE_CHECKING: from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth OPENAI_WIF_CLIENT_ID: Final = "litellm" +_IDENTITY_PARAM_ENV_NAMES: Final = ( + ("openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID"), + ("openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID"), + ("openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE"), +) _SDK_UPGRADE_MESSAGE: Final = ( "OpenAI workload identity federation requires openai>=2.32.0. " "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / " @@ -51,17 +68,25 @@ def resolve_openai_workload_identity_config( ) if static_api_key is not None: return None - effective_api_base: Final = ( - api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") - ) - if not _targets_openai_api(effective_api_base): + return _configured_openai_workload_identity(api_base=api_base, litellm_params=litellm_params) + + +def _effective_openai_api_base(api_base: str | None) -> str | None: + return api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") + + +def _configured_openai_workload_identity( + api_base: str | None, + litellm_params: Mapping[str, object] | None, +) -> OpenAIWorkloadIdentityConfig | None: + if not _targets_openai_api(_effective_openai_api_base(api_base)): return None - identity_provider_id: Final = _config_value( - litellm_params, "openai_identity_provider_id", "OPENAI_IDENTITY_PROVIDER_ID" + identity_provider_id, service_account_id, token_file = ( + _config_value(litellm_params, param_key, env_name) for param_key, env_name in _IDENTITY_PARAM_ENV_NAMES ) - service_account_id: Final = _config_value(litellm_params, "openai_service_account_id", "OPENAI_SERVICE_ACCOUNT_ID") - token_file: Final = _config_value(litellm_params, "openai_identity_token_file", "OPENAI_IDENTITY_TOKEN_FILE") if identity_provider_id is None or service_account_id is None or token_file is None: + if _carries_deployment_identity(litellm_params): + raise OpenAIError(status_code=401, message=_incomplete_identity_message(litellm_params)) return None return OpenAIWorkloadIdentityConfig( identity_provider_id=identity_provider_id, @@ -81,11 +106,200 @@ async def get_workload_identity_bearer_token_for_api_base(api_base: str) -> str return await _workload_identity_auth(config).get_token_async() +def openai_static_api_key(api_key: str | None) -> str | None: + return api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + + +def resolve_openai_bearer_token( + api_key: str | None, + api_base: str | None, + litellm_params: Mapping[str, object] | None, +) -> str | None: + static_api_key: Final = openai_static_api_key(api_key) + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=static_api_key, api_base=api_base, litellm_params=litellm_params + ) + if workload_identity_config is None: + return static_api_key + return get_workload_identity_bearer_token(workload_identity_config) + + +async def resolve_openai_bearer_token_async( + static_api_key: str | None, + api_base: str | None, + litellm_params: Mapping[str, object] | None, +) -> str | None: + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=static_api_key, api_base=api_base, litellm_params=litellm_params + ) + if workload_identity_config is None: + return static_api_key + return await _workload_identity_auth(workload_identity_config).get_token_async() + + +def _carries_deployment_identity(litellm_params: Mapping[str, object] | None) -> bool: + if litellm_params is None: + return False + return any(_param_str(litellm_params, key) is not None for key in OPENAI_WIF_KWARGS_KEYS) + + +def _incomplete_identity_message(litellm_params: Mapping[str, object] | None) -> str: + missing: Final = ", ".join( + param_key + for param_key, env_name in _IDENTITY_PARAM_ENV_NAMES + if _config_value(litellm_params, param_key, env_name) is None + ) + return ( + f"OpenAI workload identity federation on this deployment is missing {missing}. " + "Set openai_identity_provider_id, openai_service_account_id, and openai_identity_token_file " + "on the deployment, or the OPENAI_IDENTITY_PROVIDER_ID / OPENAI_SERVICE_ACCOUNT_ID / " + "OPENAI_IDENTITY_TOKEN_FILE env vars they fall back to." + ) + + +def _param_str(litellm_params: Mapping[str, object], key: str) -> str | None: + value: Final = litellm_params.get(key) + return value if isinstance(value, str) and value else None + + def _config_value(litellm_params: Mapping[str, object] | None, param_key: str, env_name: str) -> str | None: - param_value: Final = litellm_params.get(param_key) if litellm_params is not None else None - if isinstance(param_value, str) and param_value: - return param_value - return normalize_nonempty_secret_str(get_secret_str(env_name)) + param_value: Final = _param_str(litellm_params, param_key) if litellm_params is not None else None + return param_value or normalize_nonempty_secret_str(get_secret_str(env_name)) + + +def build_openai_client( + api_key: str | None, + api_base: str | None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + max_retries: int | None = None, + organization: str | None = None, + litellm_params: Mapping[str, object] | None = None, + static_key_http_client_factory: Callable[[], httpx.Client | None] | None = None, +) -> OpenAI: + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=api_key, api_base=api_base, litellm_params=litellm_params + ) + retries: Final = _max_retries_or_sdk_default(max_retries) + if workload_identity_config is None: + static_key_http_client: Final = ( + None if static_key_http_client_factory is None else static_key_http_client_factory() + ) + return OpenAI( + api_key=openai_static_api_key(api_key), + base_url=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + http_client=static_key_http_client or OpenAIHTTPClient(), + ) + cache_params: Final = _client_cache_params( + is_async=False, + workload_identity_config=workload_identity_config, + api_base=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + ) + cached: Final = BaseOpenAILLM.get_cached_openai_client( + client_initialization_params=cache_params, client_type="openai" + ) + if isinstance(cached, OpenAI): + return cached + process_http_client: Final = BaseOpenAILLM.get_sync_http_client() + client: Final = OpenAI( + workload_identity=workload_identity_config.to_sdk_workload_identity(), + base_url=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + http_client=process_http_client, + ) + BaseOpenAILLM.set_cached_openai_client( + openai_client=client, + client_type="openai", + client_initialization_params=cache_params, + litellm_owned_client=BaseOpenAILLM.owns_wrapped_http_client(process_http_client), + ) + return client + + +def build_async_openai_client( + api_key: str | None, + api_base: str | None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + max_retries: int | None = None, + organization: str | None = None, + litellm_params: Mapping[str, object] | None = None, + static_key_http_client_factory: Callable[[], httpx.AsyncClient | None] | None = None, +) -> AsyncOpenAI: + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=api_key, api_base=api_base, litellm_params=litellm_params + ) + retries: Final = _max_retries_or_sdk_default(max_retries) + if workload_identity_config is None: + static_key_http_client: Final = ( + None if static_key_http_client_factory is None else static_key_http_client_factory() + ) + return AsyncOpenAI( + api_key=openai_static_api_key(api_key), + base_url=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + http_client=static_key_http_client or OpenAIAsyncHTTPClient(), + ) + cache_params: Final = _client_cache_params( + is_async=True, + workload_identity_config=workload_identity_config, + api_base=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + ) + cached: Final = BaseOpenAILLM.get_cached_openai_client( + client_initialization_params=cache_params, client_type="openai" + ) + if isinstance(cached, AsyncOpenAI): + return cached + process_http_client: Final = BaseOpenAILLM.get_async_http_client() + client: Final = AsyncOpenAI( + workload_identity=workload_identity_config.to_sdk_workload_identity(), + base_url=api_base, + timeout=timeout, + max_retries=retries, + organization=organization, + http_client=process_http_client, + ) + BaseOpenAILLM.set_cached_openai_client( + openai_client=client, + client_type="openai", + client_initialization_params=cache_params, + litellm_owned_client=BaseOpenAILLM.owns_wrapped_http_client(process_http_client), + ) + return client + + +def _client_cache_params( + is_async: bool, + workload_identity_config: OpenAIWorkloadIdentityConfig, + api_base: str | None, + timeout: float | httpx.Timeout | None | NotGiven, + max_retries: int, + organization: str | None, +) -> dict[str, object]: + return { + "api_key": None, + "is_async": is_async, + "workload_identity_config": workload_identity_config, + "api_base": api_base, + "timeout": timeout, + "max_retries": max_retries, + "organization": organization, + } + + +def _max_retries_or_sdk_default(max_retries: int | None) -> int: + return OPENAI_SDK_DEFAULT_MAX_RETRIES if max_retries is None else max_retries def _targets_openai_api(api_base: str | None) -> bool: @@ -95,7 +309,7 @@ def _targets_openai_api(api_base: str | None) -> bool: return parsed.scheme == "https" and is_openai_backed_api_base(api_base) -@lru_cache(maxsize=16) +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _workload_identity_auth(config: OpenAIWorkloadIdentityConfig) -> WorkloadIdentityAuth: sdk_workload_identity_auth: Final = _load_sdk_workload_identity_auth() return sdk_workload_identity_auth(workload_identity=config.to_sdk_workload_identity()) diff --git a/litellm/main.py b/litellm/main.py index e6caa5c25b6..23cd3c40250 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -255,6 +255,7 @@ from .llms.openai.completion.handler import OpenAITextCompletion from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler from .llms.openai.openai import OpenAIChatCompletion from .llms.openai.transcriptions.handler import OpenAIAudioTranscription +from .llms.openai.workload_identity import build_openai_client from .llms.openai_like.chat.handler import OpenAILikeChatHandler from .llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig @@ -1212,13 +1213,18 @@ def _register_custom_pricing_for_request( litellm.register_model({shared_key: entry}, persist_across_reloads=False) return litellm.register_model( - { - deployment_id: entry, - shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry), - }, + {deployment_id: entry}, persist_across_reloads=False, warning_display_name=shared_key, custom_llm_provider=custom_llm_provider, + add_to_provider_model_sets=False, + ) + litellm.register_model( + {shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry)}, + persist_across_reloads=False, + warning_display_name=shared_key, + custom_llm_provider=custom_llm_provider, + add_to_provider_model_sets=False, ) @@ -5785,6 +5791,7 @@ def completion( "github_copilot_user_session", PROVIDER_AFFINITY_HEADER_KWARG_KEY, "fireworks_forward_user_id", + "responses_api_version", ) if key in kwargs }, @@ -7860,17 +7867,14 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompl def moderation(input: str, model: str | None = None, api_key: str | None = None, **kwargs) -> OpenAIModerationResponse: - from litellm.llms.openai.common_utils import OpenAIHTTPClient - - # only supports open ai for now - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") - - # Extract api_base from kwargs + static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") api_base: Final = kwargs.get("api_base", None) - openai_client = kwargs.get("client", None) - if openai_client is None: - openai_client = openai.OpenAI(api_key=api_key, base_url=api_base, http_client=OpenAIHTTPClient()) + openai_client: Final = kwargs.get("client", None) or build_openai_client( + api_key=static_api_key, + api_base=api_base, + litellm_params=kwargs, + ) if model is not None: response = openai_client.moderations.create(input=input, model=model) @@ -7893,8 +7897,7 @@ async def amoderation( ) -> OpenAIModerationResponse: from openai import AsyncOpenAI - # only supports open ai for now - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") optional_params: Final = GenericLiteLLMParams.model_validate(kwargs) litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) _dynamic_api_base = None @@ -7920,8 +7923,9 @@ async def amoderation( # _get_openai_client maintains in-memory caching logic for OpenAI clients _openai_client: AsyncOpenAI = openai_chat_completions.get_openai_client( is_async=True, - api_key=api_key, + api_key=static_api_key, api_base=optional_params.api_base or _dynamic_api_base, + litellm_params=optional_params.model_dump(exclude_none=True), ) else: _openai_client = openai_client @@ -7941,7 +7945,7 @@ async def amoderation( moderation_request: Final = {"input": input, "model": model} litellm_logging_obj.pre_call( input=input, - api_key=api_key, + api_key=static_api_key, additional_args={ "complete_input_dict": moderation_request, "api_base": str(_openai_client.base_url), @@ -8174,9 +8178,7 @@ def transcription( or get_secret("OPENAI_ORGANIZATION") or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 ) - # set API KEY - - api_key = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") + static_api_key: Final = api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") response = openai_audio_transcriptions.audio_transcriptions( model=model, audio_file=file, @@ -8188,7 +8190,7 @@ def transcription( logging_obj=litellm_logging_obj, max_retries=max_retries, api_base=api_base, - api_key=api_key, + api_key=static_api_key, provider_config=provider_config, litellm_params=litellm_params_dict, shared_session=shared_session, @@ -8510,6 +8512,7 @@ def speech( client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, shared_session=shared_session, + litellm_params=litellm_params_dict, ) elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 49d2ce90af4..c3efaf168ab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -342,26 +342,33 @@ }, "writer.palmyra-vision-7b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 4096, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_vision": true }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -371,13 +378,19 @@ }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -392,13 +405,17 @@ "input_cost_per_token": 1.25e-06, "input_cost_per_image_token": 1.25e-06, "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_priority": 1.75e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -489,15 +506,21 @@ "us.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 8.25e-08, "cache_read_input_token_cost_flex": 4.125e-08, + "cache_read_input_token_cost_priority": 1.44375e-07, "input_cost_per_token": 3.3e-07, + "input_cost_per_token_batches": 1.65e-07, "input_cost_per_token_flex": 1.65e-07, + "input_cost_per_token_priority": 5.775e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.75e-06, + "output_cost_per_token_batches": 1.375e-06, "output_cost_per_token_flex": 1.375e-06, + "output_cost_per_token_priority": 4.8125e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -512,13 +535,17 @@ "input_cost_per_token": 1.375e-06, "input_cost_per_image_token": 1.375e-06, "input_cost_per_audio_token": 1.375e-06, + "input_cost_per_token_batches": 6.875e-07, + "input_cost_per_token_priority": 2.40625e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5.5e-06, + "output_cost_per_token_priority": 1.925e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -548,12 +575,15 @@ "amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -561,13 +591,19 @@ }, "amazon.nova-premier-v1:0": { "cache_read_input_token_cost": 6.25e-07, + "cache_read_input_token_cost_priority": 1.09375e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_priority": 4.375e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.25e-05, + "output_cost_per_token_batches": 6.25e-06, + "output_cost_per_token_priority": 2.1875e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -578,15 +614,21 @@ "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -3987,7 +4029,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_cache_breakpoint": true, @@ -8656,7 +8698,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8705,7 +8747,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8754,7 +8796,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8803,7 +8845,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8852,7 +8894,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -23756,19 +23798,23 @@ }, "deepseek.v3.2": { "input_cost_per_token": 6.2e-07, + "input_cost_per_token_batches": 3.1e-07, "input_cost_per_token_flex": 3.1e-07, + "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 164000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, + "output_cost_per_token_batches": 9.25e-07, "output_cost_per_token_flex": 9.25e-07, + "output_cost_per_token_priority": 3.2375e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -33590,15 +33636,19 @@ "google.gemma-3-12b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 9e-08, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.9e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-12b-it.html", + "output_cost_per_token_priority": 5.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -33608,15 +33658,19 @@ "google.gemma-3-27b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.3e-07, + "input_cost_per_token_batches": 1.2e-07, "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 3.8e-07, + "output_cost_per_token_batches": 1.9e-07, "output_cost_per_token_flex": 1.9e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-27b-pt.html", + "output_cost_per_token_priority": 6.7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -33626,15 +33680,19 @@ "google.gemma-3-4b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 8e-08, + "output_cost_per_token_batches": 4e-08, "output_cost_per_token_flex": 4e-08, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-4b-it.html", + "output_cost_per_token_priority": 1.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -33747,13 +33805,19 @@ }, "global.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -40334,26 +40398,30 @@ "meta.llama3-1-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-1-8b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-8b-instruct.md", + "output_cost_per_token_batches": 1.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40408,13 +40476,15 @@ "meta.llama3-3-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-3-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40819,15 +40889,19 @@ "minimax.minimax-m2": { "deprecation_date": "2027-03-30", "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.2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2.html", + "output_cost_per_token_priority": 2.1e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -40838,25 +40912,31 @@ "minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", "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.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 196000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2-1.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "minimax.minimax-m2.5": { "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.25e-07, "output_cost_per_token": 1.2e-06, "output_cost_per_token_flex": 6e-07, "litellm_provider": "bedrock_converse", @@ -40864,10 +40944,12 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -41052,33 +41134,41 @@ "mistral.devstral-2-123b": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-07, + "input_cost_per_token_batches": 2e-07, "input_cost_per_token_flex": 2e-07, + "input_cost_per_token_priority": 7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.5e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-devstral-2-123b.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "mistral.magistral-small-2509": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 40000, "max_tokens": 40000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.63e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -41088,15 +41178,19 @@ }, "mistral.ministral-3-14b-instruct": { "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 3.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_batches": 1e-07, "output_cost_per_token_flex": 1e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41106,15 +41200,19 @@ }, "mistral.ministral-3-3b-instruct": { "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1e-07, + "output_cost_per_token_batches": 5e-08, "output_cost_per_token_flex": 5e-08, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41124,15 +41222,19 @@ }, "mistral.ministral-3-8b-instruct": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.5e-07, + "output_cost_per_token_batches": 7e-08, "output_cost_per_token_flex": 7e-08, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41177,15 +41279,19 @@ }, "mistral.mistral-large-3-675b-instruct": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.63e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41221,15 +41327,19 @@ "mistral.voxtral-mini-3b-2507": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 4e-08, + "output_cost_per_token_batches": 2e-08, "output_cost_per_token_flex": 2e-08, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-voxtral-mini-3b-2507.html", + "output_cost_per_token_priority": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": true, "supports_system_messages": true, "supports_native_structured_output": true @@ -41237,15 +41347,19 @@ "mistral.voxtral-small-24b-2507": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-voxtral-small-24b-2507.html", + "output_cost_per_token_priority": 5.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": true, "supports_system_messages": true, "supports_native_structured_output": true @@ -42056,13 +42170,17 @@ "moonshot.kimi-k2-thinking": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html", + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.38e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -42073,19 +42191,23 @@ }, "moonshotai.kimi-k2.5": { "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true }, @@ -43177,15 +43299,19 @@ "nvidia.nemotron-nano-12b-v2": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 3.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -43195,15 +43321,19 @@ "nvidia.nemotron-nano-9b-v2": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.1e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.3e-07, + "output_cost_per_token_batches": 1.2e-07, "output_cost_per_token_flex": 1.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_priority": 4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -43212,18 +43342,22 @@ }, "nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.1e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, "output_cost_per_token_flex": 1.2e-07, + "output_cost_per_token_priority": 4.2e-07, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_native_structured_output": true, "supports_response_schema": true, @@ -43231,15 +43365,19 @@ }, "nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.5e-07, + "output_cost_per_token_batches": 3.25e-07, "output_cost_per_token_flex": 3.25e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.1375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -44731,6 +44869,7 @@ "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", @@ -44739,9 +44878,10 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, "output_cost_per_token_priority": 1.05e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -44754,6 +44894,7 @@ "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, "input_cost_per_token_priority": 1.225e-07, "litellm_provider": "bedrock_converse", @@ -44762,9 +44903,10 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, "output_cost_per_token_priority": 5.25e-07, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -44772,15 +44914,19 @@ }, "openai.gpt-oss-safeguard-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, @@ -44790,15 +44936,19 @@ }, "openai.gpt-oss-safeguard-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_batches": 1e-07, "output_cost_per_token_flex": 1e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, @@ -45654,13 +45804,17 @@ "qwen.qwen3-coder-30b-a3b-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 262144, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -45668,15 +45822,19 @@ }, "qwen.qwen3-32b-v1:0": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -45796,13 +45954,17 @@ }, "qwen.qwen3-vl-235b-a22b": { "input_cost_per_token": 5.3e-07, + "input_cost_per_token_batches": 2.6e-07, + "input_cost_per_token_priority": 9.3e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.66e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.33e-06, + "output_cost_per_token_priority": 4.66e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -45812,18 +45974,22 @@ }, "qwen.qwen3-coder-next": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.75e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -47842,12 +48008,15 @@ "us.amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -47858,12 +48027,15 @@ "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -47872,15 +48044,21 @@ "us.amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -48599,14 +48777,19 @@ }, "us.deepseek.v3.2": { "input_cost_per_token": 6.2e-07, + "input_cost_per_token_batches": 3.1e-07, "input_cost_per_token_flex": 3.1e-07, + "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 1.85e-06, + "output_cost_per_token_batches": 9.25e-07, "output_cost_per_token_flex": 9.25e-07, + "output_cost_per_token_priority": 3.2375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, @@ -48648,26 +48831,30 @@ "us.meta.llama3-1-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-1-8b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-8b-instruct.md", + "output_cost_per_token_batches": 1.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -48722,13 +48909,15 @@ "us.meta.llama3-3-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-3-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -53430,37 +53619,45 @@ "zai.glm-4.7": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 203000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.2e-06, + "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_flex": 1.1e-06, + "output_cost_per_token_priority": 3.85e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-zai-glm-4-7.md", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "zai.glm-5": { "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_priority": 1.75e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -53468,19 +53665,23 @@ "zai.glm-4.7-flash": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, + "input_cost_per_token_priority": 1.225e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 203000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "output_cost_per_token_priority": 7e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-zai-glm-4-7-flash.md", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -58851,7 +59052,9 @@ }, "bedrock_mantle/openai.gpt-oss-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -58859,7 +59062,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -58872,7 +59077,9 @@ }, "bedrock_mantle/openai.gpt-oss-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, + "input_cost_per_token_priority": 1.225e-07, "output_cost_per_token": 3e-07, "output_cost_per_token_flex": 1.5e-07, "litellm_provider": "bedrock_mantle", @@ -58880,7 +59087,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.5e-07, + "output_cost_per_token_priority": 5.25e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -58893,7 +59102,9 @@ }, "bedrock_mantle/openai.gpt-oss-safeguard-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -58901,7 +59112,9 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-safeguard-120b.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -58912,7 +59125,9 @@ }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.2e-07, "output_cost_per_token": 2e-07, "output_cost_per_token_flex": 1e-07, "litellm_provider": "bedrock_mantle", @@ -58920,7 +59135,9 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-safeguard-20b.html", + "output_cost_per_token_batches": 1e-07, + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59553,7 +59770,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_prompt_caching": true, "supports_prompt_cache_breakpoint": true, "supports_reasoning": true, @@ -59592,7 +59809,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_prompt_caching": true, "supports_prompt_cache_breakpoint": true, "supports_reasoning": true, @@ -59675,7 +59892,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59714,7 +59931,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59799,7 +60016,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59834,7 +60051,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59872,7 +60089,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59907,7 +60124,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -60002,7 +60219,9 @@ }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.45e-07, "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, "litellm_provider": "bedrock_mantle", @@ -60010,6 +60229,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_priority": 7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60023,7 +60245,9 @@ }, "bedrock_mantle/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.3e-07, + "input_cost_per_token_batches": 6.5e-08, "input_cost_per_token_flex": 6.5e-08, + "input_cost_per_token_priority": 2.275e-07, "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, "litellm_provider": "bedrock_mantle", @@ -60031,6 +60255,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_priority": 7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60044,7 +60271,9 @@ }, "bedrock_mantle/google.gemma-4-e2b": { "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "output_cost_per_token": 8e-08, "output_cost_per_token_flex": 4e-08, "litellm_provider": "bedrock_mantle", @@ -60052,6 +60281,9 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4e-08, + "output_cost_per_token_priority": 1.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60068,8 +60300,10 @@ "input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_mantle", @@ -60077,6 +60311,7 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", + "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.375e-06, "supported_endpoints": [ @@ -60088,7 +60323,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/xai.grok-4.6": { "use_openai_responses_path": true, @@ -67524,7 +67759,9 @@ "bedrock_mantle/deepseek.v3.1": { "deprecation_date": "2027-03-30", "input_cost_per_token": 5.8e-07, + "input_cost_per_token_batches": 2.9e-07, "input_cost_per_token_flex": 2.9e-07, + "input_cost_per_token_priority": 1.015e-06, "output_cost_per_token": 1.68e-06, "output_cost_per_token_flex": 8.4e-07, "litellm_provider": "bedrock_mantle", @@ -67532,17 +67769,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 8.4e-07, + "output_cost_per_token_priority": 2.94e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-1.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/moonshotai.kimi-k2-thinking": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_flex": 1.25e-06, "litellm_provider": "bedrock_mantle", @@ -67550,17 +67791,21 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.38e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-235b-a22b-2507": { "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "input_cost_per_token_flex": 1.1e-07, + "input_cost_per_token_priority": 3.85e-07, "output_cost_per_token": 8.8e-07, "output_cost_per_token_flex": 4.4e-07, "litellm_provider": "bedrock_mantle", @@ -67568,17 +67813,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 4.4e-07, + "output_cost_per_token_priority": 1.54e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-235b-a22b-2507.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-32b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -67586,18 +67835,22 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-coder-30b-a3b-instruct": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -67605,17 +67858,21 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-coder-480b-a35b-instruct": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4.5e-07, + "input_cost_per_token_batches": 2.25e-07, "input_cost_per_token_flex": 2.25e-07, + "input_cost_per_token_priority": 7.875e-07, "output_cost_per_token": 1.8e-06, "output_cost_per_token_flex": 9e-07, "litellm_provider": "bedrock_mantle", @@ -67623,16 +67880,20 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 9e-07, + "output_cost_per_token_priority": 3.15e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-480b-a35b-instruct.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-next-80b-a3b-instruct": { "input_cost_per_token": 1.4e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "output_cost_per_token": 1.2e-06, "output_cost_per_token_flex": 6e-07, "litellm_provider": "bedrock_mantle", @@ -67640,17 +67901,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-next-80b-a3b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-vl-235b-a22b-instruct": { "input_cost_per_token": 5.3e-07, + "input_cost_per_token_batches": 2.6e-07, "input_cost_per_token_flex": 2.6e-07, + "input_cost_per_token_priority": 9.3e-07, "output_cost_per_token": 2.66e-06, "output_cost_per_token_flex": 1.33e-06, "litellm_provider": "bedrock_mantle", @@ -67658,13 +67923,15 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 1.33e-06, + "output_cost_per_token_priority": 4.66e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-vl-235b-a22b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, @@ -75218,5 +75485,45 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "baseten/reflection/Beam-1.0": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 1.3e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.3e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "nebius/reflection/beam": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/reflection%2Fbeam", + "supports_function_calling": true, + "supports_reasoning": true } } diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py index dcf24f974fe..db622ea3cbf 100644 --- a/litellm/proxy/_experimental/mcp_server/subscriptions.py +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -55,6 +55,7 @@ from litellm.proxy._experimental.mcp_server.catalog import ( ) from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.types.mcp import MCP_LEGACY_VERSIONS from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -300,6 +301,29 @@ async def _listing(event: ServerEvent, context: OperationContext) -> str | None: raise MCPError(code=-32603, message="Subscription catalog exceeds the pagination limit") +async def _enforce_rate_limit(context: OperationContext, server: MCPServer) -> None: + await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, server) # pyright: ignore[reportPrivateUsage] # same admission policy as other gateway operations + + +async def _unmetered(_: OperationContext, __: MCPServer) -> None: + return None + + +async def _admitted() -> ProxyRateLimitError | None: + return None + + +@asynccontextmanager +async def _listener_admission(targets: tuple[SubscriptionTarget, ...]) -> AsyncGenerator[None]: + scope: Final = operations._mcp_server_admission_memo.set( # pyright: ignore[reportPrivateUsage] # the listener's own catalog refreshes reuse its admission + {target.server.server_id: asyncio.create_task(_admitted()) for target in targets} + ) + try: + yield + finally: + operations._mcp_server_admission_memo.reset(scope) # pyright: ignore[reportPrivateUsage] # restore the caller's admission scope + + async def _recheck_authorization(check: Callable[[], Awaitable[OperationContext]]) -> None: while True: await asyncio.sleep(15) @@ -361,7 +385,7 @@ class GatewaySubscriptions: if ctx.protocol_version != "2026-07-28": raise MCPError(code=-32601, message="Subscriptions require the modern MCP protocol") async with self._capacity(owner): - return await self._listen(ctx, params, authorize) + return await self._listen(ctx, params, authorize, _enforce_rate_limit) @asynccontextmanager async def _capacity(self, owner: str) -> AsyncGenerator[None]: @@ -414,6 +438,7 @@ class GatewaySubscriptions: ctx, SubscriptionFilter(tools_list_changed=True, prompts_list_changed=True, resources_list_changed=True), authorization, + _unmetered, ) if (ready := self._initializing.get(key)) is not None: ready.set() @@ -427,6 +452,7 @@ class GatewaySubscriptions: ctx: ServerRequestContext, notifications: SubscriptionFilter, authorization: SubscriptionAuthorization, + admit: Callable[[OperationContext, MCPServer], Awaitable[None]], scope: anyio.CancelScope | None = None, *, task_status: TaskStatus[None], @@ -443,7 +469,7 @@ class GatewaySubscriptions: async with self._capacity(authorization.owner): await self._listen( - ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, ready + ctx, SubscriptionsListenRequestParams(notifications=notifications), authorization, admit, ready ) async def legacy_resource(self, ctx: ServerRequestContext, uri: str, *, subscribe: bool) -> None: @@ -464,6 +490,7 @@ class GatewaySubscriptions: ctx, SubscriptionFilter(resource_subscriptions=[uri]), state.authorization, + _enforce_rate_limit, scope, ) state.resources = {**state.resources, uri: scope} @@ -476,6 +503,7 @@ class GatewaySubscriptions: ctx: ServerRequestContext, params: SubscriptionsListenRequestParams, authorize: Callable[[], Awaitable[OperationContext]], + admit: Callable[[OperationContext, MCPServer], Awaitable[None]], ready: Callable[[SubscriptionFilter], None] | None = None, ) -> SubscriptionsListenResult: context: Final = await authorize() @@ -501,7 +529,7 @@ class GatewaySubscriptions: return fresh async def relay(target: SubscriptionTarget, *, task_status: TaskStatus[SubscriptionFilter]) -> None: - await operations._enforce_mcp_server_rate_limit(context.user_api_key_auth, target.server) # pyright: ignore[reportPrivateUsage] # same admission policy as other gateway operations + await admit(context, target.server) legacy: Final = LegacyNotifications() @@ -567,7 +595,7 @@ class GatewaySubscriptions: finally: self._upstream_slots.release() - async with anyio.create_task_group() as group: + async with _listener_admission(targets), anyio.create_task_group() as group: honored: Final = tuple( [SubscriptionFilter.model_validate(await group.start(relay, target)) for target in targets] ) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 467cb2d624c..961d32ab664 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14261,6 +14261,19 @@ "description": "Whether to fail the request if the guardrail encounters an error. Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. True (default) raises the error. False logs a critical error and lets the request proceed, so only a valid guardrail response can block or modify it.", "title": "Fail On Error" }, + "forward_messages_as_sent": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "description": "Only honored by guardrail='generic_guardrail_api'. When True, structured_messages reaches the guardrail exactly as the client sent it, so rows can carry provider-native content blocks and keys the OpenAI chat types do not declare. Default False sends OpenAI-format chat rows with undeclared keys removed.", + "title": "Forward Messages As Sent" + }, "guard_name": { "anyOf": [ { @@ -16612,6 +16625,19 @@ "description": "HTTP timeout in seconds for checking attached files. Default: 10.", "title": "File Guardrail Timeout" }, + "forward_messages_as_sent": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "description": "Only honored by guardrail='generic_guardrail_api'. When True, structured_messages reaches the guardrail exactly as the client sent it, so rows can carry provider-native content blocks and keys the OpenAI chat types do not declare. Default False sends OpenAI-format chat rows with undeclared keys removed.", + "title": "Forward Messages As Sent" + }, "gateway_name": { "anyOf": [ { diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a11268ad432..1fa276d9ab2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -6,9 +6,9 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import JSONResponse -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.pass_through.context_management import ( AnthropicContextManagementError, ) @@ -32,6 +33,7 @@ from litellm.proxy.common_request_processing import ( log_llm_api_exception, proxy_exception_from_http_exception, proxy_exception_from_route_error, + resolve_counted_tool_fields, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -47,6 +49,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli +from litellm.router import Router router: Final = APIRouter() @@ -91,6 +94,31 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO return JSONResponse(status_code=status_code, content=content, headers=exc.headers) +_TOKEN_COUNT_FIELDS: Final = ("model", "messages", "tools", "system") + + +def _token_count_request(fields: Mapping[str, object]) -> TokenCountRequest | ProxyException: + try: + return TokenCountRequest.model_validate({name: fields[name] for name in _TOKEN_COUNT_FIELDS if name in fields}) + except ValidationError as e: + return ProxyException( + message=first_validation_error(e), + type="invalid_request_error", + param=str(e.errors(include_url=False)[0]["loc"][0]), + code=status.HTTP_400_BAD_REQUEST, + ) + + +def _token_count_fields( + data: Mapping[str, object], auto_drop_params: bool, llm_router: Router | None, team_id: str | None +) -> Mapping[str, object] | ProxyException: + fields: Final = { + name: data[name] for name in (*_TOKEN_COUNT_FIELDS, "drop_params", "additional_drop_params") if name in data + } + resolving: Final = {**fields, "drop_params": True} if auto_drop_params else fields + return resolve_counted_tool_fields(resolving, llm_router, team_id) + + def _strip_total_tokens_from_anthropic_response(response: object) -> None: """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM injects into Anthropic /v1/messages responses. @@ -351,7 +379,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import count_request_tokens, proxy_config + from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -368,26 +396,21 @@ async def count_tokens( if not messages: raise HTTPException(status_code=400, detail="messages parameter is required") - # Create TokenCountRequest for the internal endpoint - from litellm.proxy._types import TokenCountRequest - - token_request: Final = TokenCountRequest( - model=model_name, - messages=messages, - tools=data.get("tools"), - system=data.get("system"), + auto_drop_params: Final = should_auto_drop_params_for_agentic_cli( + request.headers.get("user-agent", ""), "drop_params" in data, proxy_config ) + token_fields: Final = _token_count_fields(data, auto_drop_params, llm_router, user_api_key_dict.team_id) + token_request: Final = ( + token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields) + ) + if isinstance(token_request, ProxyException): + return _anthropic_error_json_response(token_request, request) token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", - **_count_tokens_body_options( - data, - should_auto_drop_params_for_agentic_cli( - request.headers.get("user-agent", ""), "drop_params" in data, proxy_config - ), - ), + **_count_tokens_body_options(data, auto_drop_params), ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 98bb0985f4a..94cf7ff4dbe 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -551,6 +551,29 @@ def _iter_fallback_targets(value: object, depth: int) -> Iterator[str | Mapping[ yield from _iter_fallback_targets(target_list, depth + 1) +def _reject_client_chosen_federation(litellm_params: Mapping[str, object], *, manages_deployments: bool) -> None: + reject_server_owned_wif_params(litellm_params) + if not manages_deployments: + reject_federated_credential_reference(litellm_params) + + +def _iter_user_config_litellm_params(user_config: object) -> Iterator[Mapping[str, object]]: + config: Final = _coerce_metadata_to_dict(user_config) + if config is None: + return + if (defaults := _coerce_metadata_to_dict(config.get("default_litellm_params"))) is not None: + yield defaults + model_list: Final = config.get("model_list") + if not isinstance(model_list, list): + return + for deployment in model_list: + if ( + isinstance(deployment, Mapping) + and (params := _coerce_metadata_to_dict(deployment.get("litellm_params"))) is not None + ): + yield params + + def iter_request_fallback_targets(request_body: Mapping[str, object]) -> Iterator[str | Mapping[str, object]]: for value in _iter_fallback_field_values(request_body): yield from _iter_fallback_targets(value, 0) @@ -658,11 +681,11 @@ def is_request_body_safe( _reject_url_valued_fallback_target(target_model) elif isinstance(target, str): _reject_url_valued_fallback_target(target) + for user_config_params in _iter_user_config_litellm_params(request_body.get("user_config")): + _reject_client_chosen_federation(user_config_params, manages_deployments=manages_deployments) litellm_params: Final = _coerce_metadata_to_dict(request_body.get("litellm_params")) if litellm_params is not None: - reject_server_owned_wif_params(litellm_params) - if not manages_deployments: - reject_federated_credential_reference(litellm_params) + _reject_client_chosen_federation(litellm_params, manages_deployments=manages_deployments) litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata")) if litellm_params_metadata is not None: _check_banned_params( diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b70cdb08b96..8a7b6dab419 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -332,7 +332,7 @@ def get_known_models_from_wildcard(wildcard_model: str, litellm_params: LiteLLM_ else: model = f"{wildcard_provider_prefix}/{model}" suffix_appended_wildcard_models.append(model) - return suffix_appended_wildcard_models or [] + return list(dict.fromkeys(suffix_appended_wildcard_models)) def expand_wildcard_deployments_for_model_info( diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 8bf39b7a17a..ead1cd8e2ca 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -17,6 +17,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param +from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit @@ -171,6 +172,8 @@ def _raise_not_found_when_openai_fallback_unservable( return if data.get("api_key") or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY"): return + if resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=data) is not None: + return raise ProxyException( message=not_found_message, type="invalid_request_error", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5f798fce6cf..2d387b43aa8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -80,6 +80,7 @@ from litellm.litellm_core_utils.served_output_texts import ( from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) +from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, @@ -129,6 +130,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ) from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError @@ -633,6 +635,63 @@ def serialize_http_exception_detail( return str(detail), None +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) +_UNSEEN_DEPLOYMENT: Final = DropPolicy(drop_params=True, additional_drop_params=frozenset()) + + +def _is_auto_router(params: Mapping[str, object]) -> bool: + model: Final = params.get("model") + return isinstance(model, str) and model.startswith(AUTO_ROUTER_MODEL_PREFIX) + + +def _named_deployments(llm_router: Router, model: str, team_id: str | None) -> tuple[Mapping[str, object], ...]: + by_id: Final = None if model in llm_router.model_names else llm_router.get_deployment(model_id=model) + if by_id is not None: + return (_REQUEST_FIELDS.validate_python(by_id.litellm_params.model_dump(exclude_none=True)),) + group: Final = llm_router.get_model_from_alias(model=model) or model + if group not in llm_router.model_names or llm_router.get_model_list_from_routing_groups(model_name=group): + return () + deployments: Final = tuple( + _REQUEST_FIELDS.validate_python(deployment["litellm_params"]) + for deployment in llm_router.get_all_deployments(model_name=group, team_id=team_id) + ) + return () if any(_is_auto_router(params) for params in deployments) else deployments + + +def _router_defaults(llm_router: Router | None) -> Mapping[str, object]: + configured: Final[object] = getattr(llm_router, "default_litellm_params", None) + if configured is None: + return {} + return {name: value for name, value in _REQUEST_FIELDS.validate_python(configured).items() if value is not None} + + +def _counted_drop_policies( + data: Mapping[str, object], llm_router: Router | None, team_id: str | None +) -> tuple[DropPolicy, ...]: + requested: Final = {**_router_defaults(llm_router), **data} + model: Final = data.get("model") + deployments: Final = ( + _named_deployments(llm_router, model, team_id) if llm_router is not None and isinstance(model, str) else () + ) + if not deployments: + return (drop_policy(requested), _UNSEEN_DEPLOYMENT) + return tuple(drop_policy({**deployment, **requested}) for deployment in deployments) + + +def resolve_counted_tool_fields( + data: Mapping[str, object], llm_router: Router | None, team_id: str | None +) -> Mapping[str, object] | ProxyException: + resolved: Final = resolve_tool_fields(data, lambda: _counted_drop_policies(data, llm_router, team_id)) + if isinstance(resolved, MalformedField): + return ProxyException( + message=resolved.message, + type="invalid_request_error", + param=resolved.field, + code=status.HTTP_400_BAD_REQUEST, + ) + return resolved + + def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: raw_detail: Final = attribute_of(exc, "detail", str(exc)) message, structured_fields = serialize_http_exception_detail(raw_detail) diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 2b8535cb113..e88b9fc2bfb 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -29,6 +29,7 @@ class SpendUpdateQueue(BaseUpdateQueue): ) -> DBSpendUpdateTransactions: """Flush all updates from the queue and return all updates aggregated by entity type.""" updates: Final = await self.flush_all_updates_from_in_memory_queue() + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) if len(updates) > 0: verbose_proxy_logger.info( "Spend tracking - flushed %d spend update items from in-memory queue", @@ -48,6 +49,7 @@ class SpendUpdateQueue(BaseUpdateQueue): "Spend update queue is full. Aggregating all entries in queue to concatenate entries." ) await self.aggregate_queue_updates() + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) async def aggregate_queue_updates(self): """Concatenate all updates in the queue to reduce the size of in-memory queue""" @@ -55,6 +57,7 @@ class SpendUpdateQueue(BaseUpdateQueue): aggregated_updates: Final = self._get_aggregated_spend_update_queue_item(updates) for update in aggregated_updates: await self.update_queue.put(update) + await self._emit_new_item_added_to_queue_event(queue_size=self.update_queue.qsize()) def _get_aggregated_spend_update_queue_item( self, updates: list[SpendUpdateQueueItem] diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 320d687b3aa..242354ac417 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -32,6 +32,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" additional_provider_specific_params=getattr(litellm_params, "additional_provider_specific_params", {}), unreachable_fallback=getattr(litellm_params, "unreachable_fallback", "fail_closed"), fail_on_error=getattr(litellm_params, "fail_on_error", True), + forward_messages_as_sent=getattr(litellm_params, "forward_messages_as_sent", False), extra_headers=getattr(litellm_params, "extra_headers", None), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 25df7a20e14..93b08831d84 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -7,10 +7,12 @@ import fnmatch import os -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Literal, Optional import httpx +from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic_core import PydanticSerializationError from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version @@ -38,6 +40,10 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_NAME: Final = "generic_guardrail_api" +_CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues) +_SENT_ROW: Final = TypeAdapter(Mapping[str, object]) +_SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]]) +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) # Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*). _HEADER_VALUE_ALLOWLIST: Final = frozenset( @@ -150,23 +156,65 @@ def _extract_inbound_headers( return None +def _is_contained_in(part: JsonValue, whole: JsonValue) -> bool: + pending: Final[list[tuple[JsonValue, JsonValue]]] = [(part, whole)] # mutable-ok: worklist, recursion is banned + while pending: + match pending.pop(): + case dict() as part_dict, dict() as whole_dict: + if not part_dict.keys() <= whole_dict.keys(): + return False + pending.extend((value, whole_dict[key]) for key, value in part_dict.items()) + case list() as part_list, list() as whole_list: + if len(part_list) != len(whole_list): + return False + pending.extend(zip(part_list, whole_list, strict=True)) + case part_leaf, whole_leaf if part_leaf != whole_leaf: + return False + case _: + pass + return True + + +def _is_shown_row(shown: AllMessageValues, returned: AllMessageValues) -> bool: + return returned == shown + + +def _is_echo(original: AllMessageValues, returned: AllMessageValues) -> bool: + """A server that parses rows with its own message types and echoes them back drops the + keys those types do not declare. That is still an echo as long as nothing changed and + every key the chat completion types declare is kept.""" + if returned == original: + return True + try: + sent: Final[JsonValue] = _JSON.validate_json(_SENT_ROW.dump_json(original)) + declared: Final[JsonValue] = _JSON.validate_json( + _CHAT_ROW.dump_json(_CHAT_ROW.validate_python(original), warnings="error") + ) + echoed: Final[JsonValue] = _JSON.validate_python(returned) + except (ValidationError, PydanticSerializationError): + return False + return _is_contained_in(echoed, sent) and _is_contained_in(declared, echoed) + + def _structured_rows_to_write_back( original_rows: Sequence[AllMessageValues] | None, - shown_rows: Sequence[AllMessageValues] | None, + compared_rows: Sequence[AllMessageValues] | None, returned_rows: Sequence[AllMessageValues], + is_echo: Callable[[AllMessageValues, AllMessageValues], bool], ) -> tuple[AllMessageValues, ...] | None: - """The request model drops row keys its message types do not declare, so a - row the server echoes back verbatim is restored to the original row object. - A server that echoes every row back unchanged has not rewritten anything - per row, so its answer is read from texts, as it was before rows could be - returned at all.""" - if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): + """A row the server echoes back is restored to the original row object. A server + that echoes every row back has not rewritten anything per row, so its answer is + read from texts, as it was before rows could be returned at all.""" + if original_rows is None or compared_rows is None or len(returned_rows) != len(original_rows): return tuple(returned_rows) - if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)): + echoed: Final = tuple( + is_echo(compared, returned) for compared, returned in zip(compared_rows, returned_rows, strict=True) + ) + if all(echoed): return None return tuple( - original if returned == shown else returned - for original, shown, returned in zip(original_rows, shown_rows, returned_rows) + original if was_echoed else returned + for original, returned, was_echoed in zip(original_rows, returned_rows, echoed, strict=True) ) @@ -200,6 +248,7 @@ class GenericGuardrailAPI(CustomGuardrail): additional_provider_specific_params: Mapping[str, object] | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", fail_on_error: bool | None = True, + forward_messages_as_sent: bool | None = False, extra_headers: list | None = None, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, @@ -235,6 +284,7 @@ class GenericGuardrailAPI(CustomGuardrail): self.unreachable_fallback: Literal["fail_closed", "fail_open"] = unreachable_fallback self.fail_on_error: bool = True if fail_on_error is None else fail_on_error + self.forward_messages_as_sent: bool = bool(forward_messages_as_sent) # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook # via getattr(guardrail_to_apply, "streaming_*", default). @@ -316,15 +366,17 @@ class GenericGuardrailAPI(CustomGuardrail): http_status_code: int | None = None, ) -> GenericGuardrailAPIInputs: status_suffix: Final = f" http_status_code={http_status_code}" if http_status_code else "" + error_type: Final = "HTTPStatusError" if isinstance(error, httpx.HTTPStatusError) else type(error).__name__ verbose_proxy_logger.critical( "Generic Guardrail API error (fail-open). Proceeding without guardrail.%s " - "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s", + "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s error_type=%s", status_suffix, getattr(self, "guardrail_name", None), getattr(self, "api_base", None), input_type, getattr(logging_obj, "litellm_call_id", None) if logging_obj else None, getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None, + error_type, exc_info=error, ) # Keep flow going - treat as action=NONE (no modifications) @@ -339,10 +391,23 @@ class GenericGuardrailAPI(CustomGuardrail): headers.update(self.headers) return headers + def _request_body( + self, + guardrail_request: GenericGuardrailAPIRequest, + structured_messages: Sequence[AllMessageValues] | None, + ) -> Mapping[str, object]: + body: Final[Mapping[str, object]] = guardrail_request.model_dump(mode="json") + if not self.forward_messages_as_sent or structured_messages is None: + return body + return { + **body, + "structured_messages": _SENT_ROWS.dump_python(_SENT_ROWS.validate_python(structured_messages), mode="json"), + } + def _build_guardrail_return_inputs( self, *, - texts: list, + texts: list[str], images: list[str] | None, tools: list[ChatCompletionToolParam] | None, structured_messages: Sequence[AllMessageValues] | None, @@ -361,8 +426,13 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + compared_rows, is_echo = ( + (structured_messages, _is_echo) if self.forward_messages_as_sent else (shown_messages, _is_shown_row) + ) rows_to_write_back: Final = ( - _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages) + _structured_rows_to_write_back( + structured_messages, compared_rows, guardrail_response.structured_messages, is_echo + ) if guardrail_response.structured_messages else None ) @@ -465,7 +535,7 @@ class GenericGuardrailAPI(CustomGuardrail): litellm_version=litellm_version, images=images, tools=([GuardrailToolParam.model_validate(t) for t in tools] if tools else None), - structured_messages=structured_messages, + structured_messages=None if self.forward_messages_as_sent else structured_messages, tool_calls=tool_calls, additional_provider_specific_params=additional_params, input_type=input_type, @@ -478,7 +548,7 @@ class GenericGuardrailAPI(CustomGuardrail): # Use mode="json" to ensure all iterables are converted to lists response: Final = await self.async_handler.post( url=self.api_base, - json=guardrail_request.model_dump(mode="json"), + json=dict(self._request_body(guardrail_request, structured_messages)), headers=headers, timeout=self.timeout, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index 74bba5c33da..fe526b72ac3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -5,6 +5,7 @@ from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams @@ -19,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" openai_moderation_guardrail: Final = OpenAIModerationGuardrail( guardrail_name=guardrail_name, + workload_identity_params=litellm_params.model_dump(include=set(OPENAI_WIF_KWARGS_KEYS), exclude_none=True), **{ **litellm_params.model_dump(exclude_none=True), "api_key": litellm_params.api_key, diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index b246b125909..33884fda4ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -3,6 +3,8 @@ OpenAI Moderation Guardrail Integration for LiteLLM """ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -17,7 +19,14 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + resolve_openai_bearer_token_async, + resolve_openai_workload_identity_config, +) from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id +from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations from litellm.types.utils import ( GenericGuardrailAPIInputs, @@ -32,6 +41,21 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +def _config_secret(value: object) -> object: + return get_secret_str(value) if isinstance(value, str) and value.startswith("os.environ/") else value + + +def _guardrail_workload_identity( + api_base: str, workload_identity_params: Mapping[str, object] +) -> OpenAIWorkloadIdentityConfig | None: + try: + return resolve_openai_workload_identity_config( + api_key=None, api_base=api_base, litellm_params=workload_identity_params + ) + except OpenAIError as incomplete_identity: + raise ValueError(f"OpenAI Moderation: {incomplete_identity.message}") from incomplete_identity + + class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): """ LiteLLM Built-in Guardrail for OpenAI Content Moderation. @@ -55,6 +79,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): model: Literal["omni-moderation-latest", "text-moderation-latest"] | None = None, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, + workload_identity_params: Mapping[str, object] | None = None, **kwargs, ): """Initialize OpenAI Moderation guardrail handler.""" @@ -78,7 +103,11 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): ) self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate - if not self.api_key: + self.workload_identity_params: Mapping[str, object] = MappingProxyType( + {key: _config_secret(value) for key, value in (workload_identity_params or {}).items()} + ) + + if not self.api_key and _guardrail_workload_identity(self.api_base, self.workload_identity_params) is None: raise ValueError( "OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration." ) @@ -109,10 +138,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): verbose_proxy_logger.debug("OpenAI Moderation guard request: %s", request_body) + bearer_token: Final = await resolve_openai_bearer_token_async( + static_api_key=self.api_key, api_base=self.api_base, litellm_params=self.workload_identity_params + ) response: Final = await self.async_handler.post( url=f"{self.api_base}/moderations", headers={ - "Authorization": f"Bearer {self.api_key}", + "Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json", }, json=request_body, diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index 72e85f35f7d..526f58dcd36 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -6,23 +6,23 @@ knowledge lives in one place. Uses guardrail translation handlers where availabl with standalone extractors for generate_content and MCP. """ +from collections.abc import Mapping from typing import Any, Final from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import list_field from litellm.types.utils import CallTypes # Call types that have no guardrail translation handler; we use standalone extractors STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {} -def _extract_generate_content_tool_names(data: dict) -> list[str]: +def _extract_generate_content_tool_names(data: Mapping[str, object]) -> list[str]: """Google generateContent: tools[].functionDeclarations[].name""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: - if not isinstance(tool, dict): - continue - for decl in tool.get("functionDeclarations") or []: + for tool in list_field(data, "tools"): + for decl in list_field(tool, "functionDeclarations"): if isinstance(decl, dict) and decl.get("name"): names.append(str(decl["name"])) return names diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 63cc8f0e3ba..b563dbc5c4c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -425,7 +425,7 @@ async def add_new_user_to_default_team( user_id: str, user_email: str | None, user_api_key_dict: UserAPIKeyAuth, - teams: list[str] | list[NewUserRequestTeam], + teams: Sequence[str | NewUserRequestTeam], prisma_client: "PrismaClient", ): tasks: Final[list[Awaitable[object]]] = [] diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 36071bfe51a..7460522f53e 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -28,7 +28,7 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.constants import LITELLM_PROXY_ADMIN_NAME, TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.ptu_pricing import ( @@ -1881,7 +1881,7 @@ async def delete_team_models( """ deleted_model_ids: Final[list[str]] = [] transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) - async with transaction_client.tx() as tx: + async with transaction_client.tx(timeout=datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx: for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ba30ac8609b..0f40662ad83 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -17,7 +17,7 @@ import traceback from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType from typing import ( @@ -42,6 +42,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.caching.redis_batch import active_request_redis_batches +from litellm.constants import TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -4674,7 +4675,7 @@ async def delete_team( # or is still waiting on the lock, in which case its own re-read happens after this commits # and sees the row gone before it writes anything. delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}} - async with prisma_client.tx() as tx: + async with prisma_client.tx(timeout=timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx: for team_id in sorted(data.team_ids): await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) await tx.litellm_teamtable.delete_many(where=delete_filter) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index e6d93aa6260..4cc849762aa 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -89,6 +89,7 @@ from litellm.proxy._types import ( Member, NewTeamRequest, NewUserRequest, + NewUserRequestTeam, NewUserResponse, ProxyErrorTypes, ProxyException, @@ -118,7 +119,11 @@ from litellm.proxy.common_utils.html_forms.jwt_display_template import ( from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler -from litellm.proxy.management_endpoints.internal_user_endpoints import new_user +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + add_new_user_to_default_team, + check_if_default_team_set, + new_user, +) from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO, LiteLLMGoogleSSO, create_generic_sso_provider from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( id_jag_assertion_capture_gap, @@ -2825,6 +2830,52 @@ async def insert_sso_user( return response +DefaultTeamJoin: TypeAlias = Callable[ + [LiteLLM_UserTable, Sequence[str | NewUserRequestTeam], PrismaClient], Awaitable[None] +] + + +async def _join_default_teams_as_admin( + user: LiteLLM_UserTable, teams: Sequence[str | NewUserRequestTeam], prisma_client: PrismaClient +) -> None: + await add_new_user_to_default_team( + user_id=user.user_id, + user_email=user.user_email, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + teams=teams, + prisma_client=prisma_client, + ) + + +def _default_teams_not_joined(user: LiteLLM_UserTable) -> tuple[str | NewUserRequestTeam, ...]: + joined: Final = frozenset(user.teams) + return tuple( + team + for team in check_if_default_team_set() or () + if (team if isinstance(team, str) else team.team_id) not in joined + ) + + +async def _insert_sso_user_or_find_concurrent_insert( + insert: Awaitable[NewUserResponse], + user_id: str | None, + prisma_client: PrismaClient, + join_default_teams: DefaultTeamJoin = _join_default_teams_as_admin, +) -> NewUserResponse | LiteLLM_UserTable: + try: + return await insert + except Exception: + users: Final = UserRepository(prisma_client, use_writer=True) + concurrent_insert: Final = None if user_id is None else await users.find_by_id(user_id) + if concurrent_insert is None: + raise + unjoined: Final = _default_teams_not_joined(concurrent_insert) + if not unjoined: + return concurrent_insert + await join_default_teams(concurrent_insert, unjoined, prisma_client) + return await users.find_by_id(concurrent_insert.user_id) or concurrent_insert + + @router.get( "/sso/get/ui_settings", tags=["experimental"], @@ -3471,9 +3522,10 @@ class SSOAuthenticationHandler: else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB - user_info = await insert_sso_user( - result_openid=result, - user_defined_values=user_defined_values, + user_info = await _insert_sso_user_or_find_concurrent_insert( + insert=insert_sso_user(result_openid=result, user_defined_values=user_defined_values), + user_id=None if user_defined_values is None else user_defined_values.get("user_id"), + prisma_client=prisma_client, ) return user_info except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 23ef89b3c5c..a7f3a4f84a7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -10,12 +10,14 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) from litellm.llms.anthropic import get_anthropic_config +from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_complete_streaming_response_legacy, @@ -254,6 +256,26 @@ class AnthropicPassthroughLoggingHandler: ), ) + @staticmethod + async def log_error_frame_as_failure( + attempt_logging_obj: LiteLLMLoggingObj, + request_body: Mapping[str, object], + all_chunks: Sequence[str | bytes], + error_event: tuple[str, str, int], + ) -> None: + """A stream closed at the provider's `event: error` frame is one failed attempt of the request: it logs + through the attempt-scoped copy its caller took at the close, so the usage it consumed never reaches + the success log of the attempt the router opens in its place.""" + error_type, message, status_code = error_event + await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( + litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks + ) + await attempt_logging_obj.dispatch_failure_handlers( + anthropic_error_frame_exception(error_type, message, status_code, str(attempt_logging_obj.model)), + "", + prefer_async_handlers=True, + ) + @staticmethod def _cost_partial_stream_or_zero( partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8f6a688fecd..6799b82bf06 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -165,6 +165,7 @@ class PassThroughStreamingHandler: ) logging_scheduled = False + stream_ended = False # rebind-ok: flipped once the upstream ran to its end, so the finally can tell a close-at-frame from a completed stream model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( request_body=request_body, url_route=url_route, @@ -224,6 +225,7 @@ class PassThroughStreamingHandler: ) if pending: yield pending + stream_ended = True # rebind-ok: the finally reads it to spot a close-at-frame # Stream completed cleanly. When the proxy armed deferred # dispatch (post-call guardrails active), park the logging # coroutine on logging_obj instead of enqueueing now, so @@ -267,13 +269,46 @@ class PassThroughStreamingHandler: # a success would double-log the same request. if not logging_scheduled and raw_bytes and response.status_code < 400: logging_scheduled = True + superseded_error_event: Final = ( + None + if stream_ended + else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes) + ) try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=( + _build_logging_coroutine() + if superseded_error_event is None + else AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( + attempt_logging_obj=litellm_logging_obj.attempt_scoped_copy(), + request_body=resolved_request_body, + all_chunks=tuple(raw_bytes), + error_event=superseded_error_event, + ) + ) + ) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) else: bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) + @staticmethod + def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: + """The provider `event: error` frame of a stream the router closed because a retry or fallback took the + request over. That attempt is logged as its own failure, on a copy taken now so the retry's rebinding of + the shared logging object never reaches it. A stream that ran to its end or that the client closed keeps + logging through the success route, so its usage is still billed.""" + from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( + closing_superseded_attempt, + parse_anthropic_error_event, + ) + + if endpoint_type != EndpointType.ANTHROPIC or not closing_superseded_attempt(): + return None + + complete_frames, _ = split_complete_sse_frames(b"".join(raw_bytes)) + return parse_anthropic_error_event(complete_frames) + @staticmethod async def route_streaming_logging_to_handler( litellm_logging_obj: LiteLLMLoggingObj, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index d9f5dd9c39d..778d40c2d07 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -34,6 +34,7 @@ from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, build_blocked_response, @@ -45,7 +46,12 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket, ) -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, cost_headers, create_response +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + cost_headers, + create_response, + resolve_counted_tool_fields, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -58,6 +64,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import handle_exception_on_proxy +from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -70,7 +77,6 @@ from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging - from litellm.router import Router router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) @@ -1377,6 +1383,15 @@ def _proxy_token_counter() -> _TokenCounter: _token_counter_dep: Final = Depends(_proxy_token_counter) +def _proxy_llm_router() -> Router | None: + from litellm.proxy.proxy_server import llm_router + + return llm_router + + +_llm_router_dep: Final = Depends(_proxy_llm_router) + + def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: body: Final[_ResponsesApiErrorBody] = { "error": { @@ -1433,6 +1448,8 @@ def _responses_input_as_token_count_messages( async def responses_input_tokens( request: Request, token_counter: _TokenCounter = _token_counter_dep, + llm_router: Router | None = _llm_router_dep, + user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, ): """ Count the input tokens of a Responses API request without calling the model. @@ -1465,6 +1482,13 @@ async def responses_input_tokens( code="missing_required_parameter", ) + counted_fields: Final = { + name: data[name] for name in ("model", "tools", "drop_params", "additional_drop_params") if name in data + } + token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) + if isinstance(token_fields, ProxyException): + return _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + try: payload: Final[_TokenCountPayload] = { "model": model_name, @@ -1472,9 +1496,13 @@ async def responses_input_tokens( input_value=input_value, instructions=data.get("instructions"), ), - "tools": data.get("tools"), + "tools": token_fields.get("tools"), } token_request: Final = TokenCountRequest.model_validate(payload) + except ValidationError as e: + return _responses_invalid_request_response( + message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + ) except Exception as e: return _responses_invalid_request_response( message=f"Invalid request for token counting: {e}", param=None, code=None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index dc6423c979b..36ace3c1697 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3493,6 +3493,7 @@ class ProxyLogging: Is triggered when self._is_proxy_only_error() returns True """ litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) + set_up_by_request: Final = litellm_logging_obj is not None if litellm_logging_obj is None: from litellm._uuid import uuid @@ -3557,7 +3558,9 @@ class ProxyLogging: input = request_data["input"] litellm_logging_obj.model_call_details["input"] = input body_shape_call_type = CallTypes.aembedding.value - resolved_call_type: Final = _call_type_for_route(route) or body_shape_call_type + resolved_call_type: Final = _call_type_for_route(route) or ( + None if set_up_by_request else body_shape_call_type + ) if resolved_call_type is not None and litellm_logging_obj.call_type != CallTypes.pass_through.value: litellm_logging_obj.call_type = resolved_call_type litellm_logging_obj.model_call_details["call_type"] = resolved_call_type diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index aeec0bf67c0..4acce4de586 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -43,6 +43,7 @@ from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import realtime_ssl_for_url from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime +from ..llms.openai.workload_identity import openai_static_api_key, resolve_openai_bearer_token_async from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config from ..llms.vertex_ai.vertex_llm_base import VertexBase @@ -130,6 +131,17 @@ def _get_realtime_http_provider_config( return provider_config, resolved_api_base.rstrip("/"), resolved_api_key +async def _realtime_http_bearer_token( + custom_llm_provider: str, api_key: str, api_base: str, litellm_params: GenericLiteLLMParams +) -> str: + if custom_llm_provider != LlmProviders.OPENAI.value: + return api_key + bearer_token: Final = await resolve_openai_bearer_token_async( + static_api_key=api_key or None, api_base=api_base, litellm_params=litellm_params.model_dump(exclude_none=True) + ) + return bearer_token or "" + + @wrapper_client async def acreate_realtime_client_secret( model: str | None = None, @@ -180,7 +192,9 @@ async def acreate_realtime_client_secret( request_data["session"] = _with_resolved_session_model(request_data["session"], model_name) return await base_llm_http_handler.async_realtime_client_secret_handler( api_base=resolved_api_base, - api_key=resolved_api_key, + api_key=await _realtime_http_bearer_token( + custom_llm_provider, resolved_api_key, resolved_api_base, litellm_params + ), request_data=request_data, logging_obj=litellm_logging_obj, timeout=timeout or request_timeout, @@ -253,7 +267,9 @@ async def acreate_realtime_transcription_session( request_data["input_audio_transcription"]["model"] = model_name return await base_llm_http_handler.async_realtime_transcription_session_handler( api_base=resolved_api_base, - api_key=resolved_api_key, + api_key=await _realtime_http_bearer_token( + custom_llm_provider, resolved_api_key, resolved_api_base, litellm_params + ), request_data=request_data, logging_obj=litellm_logging_obj, timeout=timeout or request_timeout, @@ -418,17 +434,21 @@ async def _alive_session( api_base: str | None, api_key: str | None, litellm_metadata: Mapping[str, object], + litellm_params: Mapping[str, object], ) -> None: rejection: Final = _live_session_rejection(custom_llm_provider, litellm_metadata) if rejection is not None: _raise_live_session_rejection(rejection) + resolved_api_base: Final = api_base or litellm.api_base or "https://api.openai.com/" await openai_live_sessions.async_live_session( model=model, websocket=websocket, logging_obj=litellm_logging_obj, session_start=session_start, - api_base=api_base or litellm.api_base or "https://api.openai.com/", - api_key=api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY"), + api_base=resolved_api_base, + api_key=await resolve_openai_bearer_token_async( + static_api_key=openai_static_api_key(api_key), api_base=resolved_api_base, litellm_params=litellm_params + ), ) @@ -497,6 +517,7 @@ async def _arealtime( api_base=dynamic_api_base or litellm_params.api_base, api_key=dynamic_api_key, litellm_metadata=_build_litellm_metadata(kwargs), + litellm_params=kwargs, ) return @@ -558,8 +579,9 @@ async def _arealtime( ) elif _custom_llm_provider == "openai": api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or "https://api.openai.com/" - # set API KEY - api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") + api_key = await resolve_openai_bearer_token_async( + static_api_key=openai_static_api_key(dynamic_api_key), api_base=api_base, litellm_params=kwargs + ) await openai_realtime.async_realtime( model=model, @@ -774,9 +796,18 @@ async def realtime_health_check( str | None, api_version or resolved_params.get("api_version") ) url: str | None = None + health_check_api_key: Final = ( + await resolve_openai_bearer_token_async( + static_api_key=openai_static_api_key(resolved_api_key), + api_base=resolved_api_base, + litellm_params=resolved_params, + ) + if custom_llm_provider == LlmProviders.OPENAI.value + else resolved_api_key + ) auth_headers: Final = _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, - api_key=resolved_api_key, + api_key=health_check_api_key, model_params=resolved_params, ) if custom_llm_provider == "azure": diff --git a/litellm/router.py b/litellm/router.py index 914f5eb1271..157fb59d9ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -123,6 +123,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_structure, ) from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.llms.base_llm.vector_store.transformation import ( @@ -501,7 +502,9 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") -_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) +_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( + {"model", "api_base", "api_key", "api_version", "responses_api_version"} +) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" _CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 @@ -538,6 +541,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) +_NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) _SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -574,6 +578,21 @@ def _deployment_that_just_failed(request_metadata: object) -> str | None: return model_id if isinstance(model_id, str) else None +def _drops_malformed_request_tools( + deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] +) -> bool: + if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"): + return False + set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} + return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") + + +def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None: + if not isinstance(request_tools, (list, tuple)): + return None + return [*deployment_tools, *request_tools] if deployment_tools or request_tools else None + + def _deployment_pick_attributes(model: str, request_kwargs: Mapping[str, object] | None) -> Mapping[str, str | int]: """Bounded attributes for one deployment pick; attempt is 1-based within the current model group.""" kwargs: Final = request_kwargs or {} @@ -4508,7 +4527,9 @@ class Router: return deployment_pydantic_obj @staticmethod - def _merge_tools_from_deployment(deployment: dict, kwargs: dict) -> None: + def _merge_tools_from_deployment( + deployment: dict, kwargs: dict, defaults: Mapping[str, object] = _NO_DEFAULT_PARAMS + ) -> None: """ Merge tools from deployment litellm_params with request kwargs. When both have tools, concatenate them (deployment tools first, then request tools). @@ -4519,11 +4540,12 @@ class Router: dep_params = dep_params_raw else: dep_params = dep_params_raw.model_dump(exclude_none=True) - dep_tools: Final = dep_params.get("tools") or [] - req_tools: Final = kwargs.get("tools") or [] - if dep_tools or req_tools: - merged: Final = list(dep_tools) + list(req_tools) - kwargs["tools"] = merged + if _drops_malformed_request_tools(dep_params, defaults, kwargs): + for name in TOOL_FIELDS: + kwargs.pop(name, None) + merged_tools: Final = _merged_tools(dep_params.get("tools") or [], kwargs.get("tools") or []) + if merged_tools is not None: + kwargs["tools"] = merged_tools if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: kwargs["tool_choice"] = dep_params["tool_choice"] @@ -4544,7 +4566,7 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] @@ -6042,6 +6064,7 @@ class Router: """ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( aclose_if_supported, + aclose_superseded_attempt, parse_anthropic_error_event, parse_anthropic_refusal_stop_details, ) @@ -6051,6 +6074,14 @@ class Router: source_iterator: Final = response + async def close_source() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + + async def close_superseded() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_superseded_attempt(source_iterator) + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: from litellm.exceptions import MidStreamFallbackError @@ -6137,11 +6168,11 @@ class Router: model, initial_kwargs, wrapper, + close_superseded, ): yield item finally: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(source_iterator) + await close_source() # Referenced by stream_with_fallbacks via closure - assigned here, before # the generator body ever runs, so the reference resolves fine despite @@ -6157,6 +6188,7 @@ class Router: model: str, initial_kwargs: dict[str, Any], # mutable-ok: handed to _aanthropic_messages_fallback_attempt, which mutates it wrapper: "FallbackAwareAnthropicMessagesStream", + close_superseded: Callable[[], Awaitable[None]], ) -> AsyncGenerator[bytes, None]: """Turns a source-iterator failure into a fallback attempt or the error reaching the caller.""" from litellm.exceptions import MidStreamFallbackError @@ -6176,6 +6208,7 @@ class Router: ) if fallback_error is None: raise stream_error + await close_superseded() outcome: Final = await self._astream_retry_same_group( fallback_error, initial_kwargs, @@ -11035,6 +11068,7 @@ class Router: persist_across_reloads=False, warning_display_name=model, custom_llm_provider=custom_llm_provider, + add_to_provider_model_sets=False, ) ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index c4e4d9ccc0b..9b9a3d185cc 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -309,6 +309,9 @@ def filter_team_based_models( ] +_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) + + def _deployment_supports_web_search(deployment: dict) -> bool: """ Check if a deployment supports web search. @@ -343,14 +346,13 @@ def filter_web_search_deployments( if isinstance(healthy_deployments, dict): return healthy_deployments - is_web_search_request = False - tools: Final = request_kwargs.get("tools") or [] - for tool in tools: - # These are the two websearch tools for OpenAI / Azure. - if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview": - is_web_search_request = True - break - + tools: Final[object] = request_kwargs.get("tools") + is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( + isinstance(tool, dict) + and isinstance(tool_type := tool.get("type"), str) + and tool_type in _WEB_SEARCH_TOOL_TYPES + for tool in tools + ) if not is_web_search_request: return healthy_deployments diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c7e3dab0d1a..bb64d912c4c 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -329,6 +329,12 @@ def _should_run_cooldown_logic( ) return False + if isinstance(original_exception, litellm.MalformedRequestFieldError): + verbose_router_logger.debug( + "Should Not Run Cooldown Logic: litellm rejected the request before calling the deployment" + ) + return False + if request_kwargs is not None and _is_per_user_provider_request(request_kwargs): verbose_router_logger.debug( "Should Not Run Cooldown Logic: the failure came from one caller's per-user provider session" diff --git a/litellm/systemone/main.py b/litellm/systemone/main.py index c91f504eefa..01f151ef4d0 100644 --- a/litellm/systemone/main.py +++ b/litellm/systemone/main.py @@ -1,4 +1,6 @@ +import asyncio from collections.abc import Mapping +from functools import partial from typing import Final import httpx @@ -75,21 +77,24 @@ async def asystemone( assets: DecisionsAssetEntries | None = None, **kwargs: object, # kwargs-ok: litellm_params forwarded by @client ) -> DecisionsResponse: - call: Final = _prepare_systemone_call( - model=model, - state=state, - questions=questions, - images=images, - assets=assets, - safety_identifier=safety_identifier, - audio=audio, - videos=videos, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, + call: Final = await asyncio.to_thread( + partial( + _prepare_systemone_call, + model=model, + state=state, + questions=questions, + images=images, + assets=assets, + safety_identifier=safety_identifier, + audio=audio, + videos=videos, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) ) return _systemone_response(await asend(call)) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index f20a3fa3d7c..33db8cbc6f0 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1195,6 +1195,16 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="Maximum decisions calls in flight at once on the guardrail in each proxy worker. Implemented by guardrail='decision_model'.", ) + forward_messages_as_sent: bool | None = Field( + default=False, + description=( + "Only honored by guardrail='generic_guardrail_api'. When True, structured_messages reaches the " + "guardrail exactly as the client sent it, so rows can carry provider-native content blocks and " + "keys the OpenAI chat types do not declare. Default False sends OpenAI-format chat rows with " + "undeclared keys removed." + ), + ) + extra_headers: list[str] | None = Field( default=None, description=( diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 8387afa22bb..174d07a2cea 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -67,6 +67,7 @@ class ProviderConnection: api_key: str | None = None api_base: str | None = None api_version: str | None = None + responses_api_version: str | None = None region_name: str | None = None headers: Mapping[str, str] | None = None provider_specific_header: "ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None" = None diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 49eaa154fec..d34a3a13838 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,7 +1,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles -from pydantic import ConfigDict, Field +from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer from typing_extensions import TypedDict from litellm.types.llms.base import LiteLLMBaseModel @@ -17,12 +17,20 @@ class GuardrailToolParam(LiteLLMBaseModel): """A tool forwarded verbatim to the guardrail for inspection. Built-in tools (code_interpreter, file_search, ...) have no ``function`` block - and stash their config in tool-specific keys, so only ``type`` is required and - ``extra="allow"`` preserves the rest instead of stripping it. + and stash their config in tool-specific keys, and provider-native tools such as + Gemini ``function_declarations`` carry no ``type`` at all, so nothing is required + and ``extra="allow"`` preserves the rest instead of stripping it. """ model_config = ConfigDict(extra="allow") - type: str + type: str | None = None + + @model_serializer(mode="wrap") + def _without_absent_type(self, handler: SerializerFunctionWrapHandler) -> Mapping[str, object]: + dumped: Final = cast("Mapping[str, object]", handler(self)) # cast-ok: handler dumps this model's fields + if self.type is not None: + return dumped + return {key: value for key, value in dumped.items() if key != "type"} class GenericGuardrailAPIMetadata(TypedDict, total=False): @@ -62,6 +70,15 @@ class GenericGuardrailAPIOptionalParams(LiteLLMBaseModel): ), ) + forward_messages_as_sent: bool | None = Field( + default=False, + description=( + "When True, structured_messages reaches the guardrail exactly as the client sent it, so rows can " + "carry provider-native content blocks and keys the OpenAI chat types do not declare. False " + "(default) sends OpenAI-format chat rows with undeclared keys removed." + ), + ) + streaming_end_of_stream_only: bool | None = Field( default=None, description=( diff --git a/litellm/types/router.py b/litellm/types/router.py index 0ca39e2400d..96add8bb2a1 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -496,6 +496,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): def __init__(self, /, **data: object) -> None: ... # kwargs-ok: provider parameters vary by deployment custom_llm_provider: str | None = None + responses_api_version: str | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None @@ -703,6 +704,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): api_key: str | None api_base: str | None api_version: str | None + responses_api_version: ReadOnly[str | None] timeout: float | str | httpx.Timeout | SDKTimeout | None # writable-ok: preserve timeout updates stream_timeout: float | str | None max_retries: int | None @@ -1400,8 +1402,9 @@ def reject_server_owned_wif_params(body: Mapping[str, object]) -> None: This lives here rather than under ``litellm.proxy`` so the router can call it on a post-authentication merge without core importing from the proxy package. """ + submitted_names: Final = frozenset(key.removesuffix("[]") for key in body) for param in _server_owned_wif_litellm_params: - if param in body: + if param in submitted_names: raise ValueError( f"Rejected Request: {param} is a server-owned workload identity federation or OAuth token exchange " "parameter " diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a97e35a9021..054a0a1c034 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3900,6 +3900,7 @@ class MirroredPricingParams(LiteLLMBaseModel): output_cost_per_character: float | None = None cache_read_input_token_cost: float | None = None cache_creation_input_token_cost: float | None = None + cache_creation_input_token_cost_above_1hr: float | None = None tiered_pricing: list[dict[str, Any]] | None = None @@ -3927,7 +3928,6 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_priority: float | None = None input_cost_per_token_balanced: float | None = None input_cost_per_token_ultrafast: float | None = None - cache_creation_input_token_cost_above_1hr: float | None = None cache_creation_input_token_cost_above_100k_tokens: float | None = None cache_creation_input_token_cost_above_1hr_above_100k_tokens: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 8aa6eb23c68..3b43c570466 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -92,6 +92,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens +from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -474,6 +475,7 @@ from .exceptions import ( BudgetExceededError, ContentPolicyViolationError, ContextWindowExceededError, + MalformedRequestFieldError, ModelNotMappedError, NotFoundError, OpenAIError, @@ -1495,7 +1497,23 @@ def check_coroutine(value) -> bool: return get_coroutine_checker().is_async_callable(value) -async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str): +def _deployment_provider(kwargs: Mapping[str, object]) -> str: + model: Final = kwargs.get("model") + custom_llm_provider: Final = kwargs.get("custom_llm_provider") + api_base: Final = kwargs.get("api_base") + if not isinstance(model, str): + return "" + try: + return litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + api_base=api_base if isinstance(api_base, str) else None, + )[1] + except BadRequestError: + return "" + + +async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) -> dict[str, Any]: """ Allow modifying the request just before it's sent to the deployment. @@ -1506,7 +1524,16 @@ async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) except ValueError: typed_call_type = None # unknown call type - modified_kwargs = kwargs.copy() + route: Final[object] = getattr(kwargs.get("litellm_logging_obj"), "call_type", None) + resolved: Final = resolve_deployment_tool_fields(kwargs, route) + if isinstance(resolved, MalformedField): + raise MalformedRequestFieldError( + message=resolved.message, + model=kwargs.get("model"), + llm_provider=_deployment_provider(kwargs), + body={"type": "invalid_request_error", "param": resolved.field}, + ) + modified_kwargs = dict(resolved) custom_logger_class: Final = _get_cached_custom_logger() for callback in litellm.callbacks: @@ -3418,6 +3445,7 @@ def register_model( persist_across_reloads: bool = True, warning_display_name: str | None = None, custom_llm_provider: str | None = None, + add_to_provider_model_sets: bool = True, ): """ Register new / Override existing models (and their pricing) to specific providers. @@ -3446,6 +3474,10 @@ def register_model( ``custom_llm_provider`` scopes the built-in cost-map match to entries for that provider, so a deployment id that happens to equal another provider's catalog key stays its own provider-less entry instead of merging into it. + + ``add_to_provider_model_sets`` adds the keys to the per-provider model sets + behind provider resolution and wildcard model listing. Pass False when the + keys are opaque ids rather than model names (e.g. router deployment ids). """ loaded_model_cost = {} @@ -3531,56 +3563,10 @@ def register_model( _invalidate_model_cost_lowercase_map() verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key) - # add new model names to provider lists - if value.get("litellm_provider") == "openai": - if key not in litellm.open_ai_chat_completion_models: - litellm.open_ai_chat_completion_models.add(key) - elif value.get("litellm_provider") == "text-completion-openai": - if key not in litellm.open_ai_text_completion_models: - litellm.open_ai_text_completion_models.add(key) - elif value.get("litellm_provider") == "cohere": - if key not in litellm.cohere_models: - litellm.cohere_models.add(key) - elif value.get("litellm_provider") == "anthropic": - if key not in litellm.anthropic_models: - litellm.anthropic_models.add(key) - elif value.get("litellm_provider") == "openrouter": - split_string = key.split("/", 1) - if split_string[-1] not in litellm.openrouter_models: - litellm.openrouter_models.add(split_string[-1]) - elif value.get("litellm_provider") == "vercel_ai_gateway": - if key not in litellm.vercel_ai_gateway_models: - litellm.vercel_ai_gateway_models.add(key) - elif value.get("litellm_provider") == "edenai": - if key not in litellm.edenai_models: - litellm.edenai_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-text-models": - if key not in litellm.vertex_text_models: - litellm.vertex_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-text-models": - if key not in litellm.vertex_code_text_models: - litellm.vertex_code_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-chat-models": - if key not in litellm.vertex_chat_models: - litellm.vertex_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-chat-models": - if key not in litellm.vertex_code_chat_models: - litellm.vertex_code_chat_models.add(key) - elif value.get("litellm_provider") == "ai21": - if key not in litellm.ai21_models: - litellm.ai21_models.add(key) - elif value.get("litellm_provider") == "nlp_cloud": - if key not in litellm.nlp_cloud_models: - litellm.nlp_cloud_models.add(key) - elif value.get("litellm_provider") == "aleph_alpha": - if key not in litellm.aleph_alpha_models: - litellm.aleph_alpha_models.add(key) - elif value.get("litellm_provider") == "bedrock": - if key not in litellm.bedrock_models: - litellm.bedrock_models.add(key) - elif value.get("litellm_provider") == "novita": - if key not in litellm.novita_models: - litellm.novita_models.add(key) + if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter": + litellm.openrouter_models.add(_key_str.split("/", 1)[-1]) + if add_to_provider_model_sets: + litellm.add_known_models(loaded_model_cost) return model_cost @@ -9252,7 +9238,7 @@ class ProviderConfigManager: return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: - if "gpt-4o" in model: + if "gpt-4o" in model or (model.startswith("gpt-") and "transcribe" in model): return litellm.OpenAIGPTAudioTranscriptionConfig() else: return litellm.OpenAIWhisperAudioTranscriptionConfig() diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index e0b536d6c57..b93d9874dbc 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -1,6 +1,6 @@ # litellm/proxy/vector_stores/vector_store_registry.py import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from itertools import chain from typing import ( @@ -11,7 +11,7 @@ from typing import ( get_args, ) -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices @@ -34,6 +34,26 @@ else: PrismaClient = Any _DB_ROW: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) +_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) +_TOOL: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) +_STORE_ID_ITEMS: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) + + +def _named_store_ids(store_ids: object) -> tuple[str, ...]: + if isinstance(store_ids, str): + return (store_ids,) + try: + items: Final = _STORE_ID_ITEMS.validate_python(store_ids) + except ValidationError: + return () + return tuple(store_id for store_id in items if isinstance(store_id, str)) + + +def _tool_vector_store_ids(tool: object) -> tuple[str, ...]: + try: + return _named_store_ids(_TOOL.validate_python(tool).get("vector_store_ids")) + except ValidationError: + return () class VectorStoreIndexRegistry: @@ -131,7 +151,7 @@ class VectorStoreRegistry: return VectorStoreToolParams(**kwargs) - def get_vector_store_ids_to_run(self, non_default_params: dict, tools: list[dict] | None = None) -> list[str]: + def get_vector_store_ids_to_run(self, non_default_params: dict, tools: object = None) -> list[str]: """ Returns the vector store ids to run @@ -175,7 +195,7 @@ class VectorStoreRegistry: tools_to_remove: Final[list[int]] = [] for i, tool in enumerate(tools): - tool_vector_store_ids = tool.get("vector_store_ids", []) + tool_vector_store_ids = _tool_vector_store_ids(tool) if not tool_vector_store_ids: continue @@ -395,14 +415,15 @@ class VectorStoreRegistry: return vector_stores_to_run - def _get_vector_store_ids_from_tool_calls(self, tools: list[dict] | None, vector_store_ids: list[str]) -> list[str]: + def _get_vector_store_ids_from_tool_calls(self, tools: object, vector_store_ids: list[str]) -> list[str]: """ Returns the vector store ids from the tool calls """ - return [ - *vector_store_ids, - *chain.from_iterable(tool["vector_store_ids"] for tool in tools or () if "vector_store_ids" in tool), - ] + try: + validated_tools: Final = _ITEMS.validate_python(tools) + except ValidationError: + return vector_store_ids + return [*vector_store_ids, *chain.from_iterable(_tool_vector_store_ids(tool) for tool in validated_tools)] def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 49d2ce90af4..c3efaf168ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -342,26 +342,33 @@ }, "writer.palmyra-vision-7b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 4096, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_vision": true }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -371,13 +378,19 @@ }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -392,13 +405,17 @@ "input_cost_per_token": 1.25e-06, "input_cost_per_image_token": 1.25e-06, "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_priority": 1.75e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -489,15 +506,21 @@ "us.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 8.25e-08, "cache_read_input_token_cost_flex": 4.125e-08, + "cache_read_input_token_cost_priority": 1.44375e-07, "input_cost_per_token": 3.3e-07, + "input_cost_per_token_batches": 1.65e-07, "input_cost_per_token_flex": 1.65e-07, + "input_cost_per_token_priority": 5.775e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.75e-06, + "output_cost_per_token_batches": 1.375e-06, "output_cost_per_token_flex": 1.375e-06, + "output_cost_per_token_priority": 4.8125e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -512,13 +535,17 @@ "input_cost_per_token": 1.375e-06, "input_cost_per_image_token": 1.375e-06, "input_cost_per_audio_token": 1.375e-06, + "input_cost_per_token_batches": 6.875e-07, + "input_cost_per_token_priority": 2.40625e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5.5e-06, + "output_cost_per_token_priority": 1.925e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -548,12 +575,15 @@ "amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -561,13 +591,19 @@ }, "amazon.nova-premier-v1:0": { "cache_read_input_token_cost": 6.25e-07, + "cache_read_input_token_cost_priority": 1.09375e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_priority": 4.375e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.25e-05, + "output_cost_per_token_batches": 6.25e-06, + "output_cost_per_token_priority": 2.1875e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -578,15 +614,21 @@ "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -3987,7 +4029,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_cache_breakpoint": true, @@ -8656,7 +8698,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8705,7 +8747,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8754,7 +8796,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8803,7 +8845,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -8852,7 +8894,7 @@ "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -23756,19 +23798,23 @@ }, "deepseek.v3.2": { "input_cost_per_token": 6.2e-07, + "input_cost_per_token_batches": 3.1e-07, "input_cost_per_token_flex": 3.1e-07, + "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 164000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, + "output_cost_per_token_batches": 9.25e-07, "output_cost_per_token_flex": 9.25e-07, + "output_cost_per_token_priority": 3.2375e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -33590,15 +33636,19 @@ "google.gemma-3-12b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 9e-08, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.9e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-12b-it.html", + "output_cost_per_token_priority": 5.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -33608,15 +33658,19 @@ "google.gemma-3-27b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.3e-07, + "input_cost_per_token_batches": 1.2e-07, "input_cost_per_token_flex": 1.2e-07, + "input_cost_per_token_priority": 4e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 3.8e-07, + "output_cost_per_token_batches": 1.9e-07, "output_cost_per_token_flex": 1.9e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-27b-pt.html", + "output_cost_per_token_priority": 6.7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -33626,15 +33680,19 @@ "google.gemma-3-4b-it": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 8e-08, + "output_cost_per_token_batches": 4e-08, "output_cost_per_token_flex": 4e-08, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-google-gemma-3-4b-it.html", + "output_cost_per_token_priority": 1.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -33747,13 +33805,19 @@ }, "global.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -40334,26 +40398,30 @@ "meta.llama3-1-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-1-8b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-8b-instruct.md", + "output_cost_per_token_batches": 1.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40408,13 +40476,15 @@ "meta.llama3-3-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-3-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40819,15 +40889,19 @@ "minimax.minimax-m2": { "deprecation_date": "2027-03-30", "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.2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2.html", + "output_cost_per_token_priority": 2.1e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -40838,25 +40912,31 @@ "minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", "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.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 196000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-minimax-minimax-m2-1.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "minimax.minimax-m2.5": { "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.25e-07, "output_cost_per_token": 1.2e-06, "output_cost_per_token_flex": 6e-07, "litellm_provider": "bedrock_converse", @@ -40864,10 +40944,12 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -41052,33 +41134,41 @@ "mistral.devstral-2-123b": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-07, + "input_cost_per_token_batches": 2e-07, "input_cost_per_token_flex": 2e-07, + "input_cost_per_token_priority": 7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2e-06, + "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, + "output_cost_per_token_priority": 3.5e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-devstral-2-123b.html", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "mistral.magistral-small-2509": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 40000, "max_tokens": 40000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.63e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -41088,15 +41178,19 @@ }, "mistral.ministral-3-14b-instruct": { "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 3.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_batches": 1e-07, "output_cost_per_token_flex": 1e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41106,15 +41200,19 @@ }, "mistral.ministral-3-3b-instruct": { "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1e-07, + "output_cost_per_token_batches": 5e-08, "output_cost_per_token_flex": 5e-08, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41124,15 +41222,19 @@ }, "mistral.ministral-3-8b-instruct": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.5e-07, + "output_cost_per_token_batches": 7e-08, "output_cost_per_token_flex": 7e-08, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41177,15 +41279,19 @@ }, "mistral.mistral-large-3-675b-instruct": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.8e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 2.63e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -41221,15 +41327,19 @@ "mistral.voxtral-mini-3b-2507": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 4e-08, + "output_cost_per_token_batches": 2e-08, "output_cost_per_token_flex": 2e-08, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-voxtral-mini-3b-2507.html", + "output_cost_per_token_priority": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": true, "supports_system_messages": true, "supports_native_structured_output": true @@ -41237,15 +41347,19 @@ "mistral.voxtral-small-24b-2507": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, + "input_cost_per_token_priority": 1.7e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-mistral-ai-voxtral-small-24b-2507.html", + "output_cost_per_token_priority": 5.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": true, "supports_system_messages": true, "supports_native_structured_output": true @@ -42056,13 +42170,17 @@ "moonshot.kimi-k2-thinking": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html", + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.38e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -42073,19 +42191,23 @@ }, "moonshotai.kimi-k2.5": { "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.25e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true }, @@ -43177,15 +43299,19 @@ "nvidia.nemotron-nano-12b-v2": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 3.5e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-12b-v2-vl-bf16.html", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -43195,15 +43321,19 @@ "nvidia.nemotron-nano-9b-v2": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.1e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.3e-07, + "output_cost_per_token_batches": 1.2e-07, "output_cost_per_token_flex": 1.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-nvidia-nvidia-nemotron-nano-9b-v2.html", + "output_cost_per_token_priority": 4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_response_schema": true, @@ -43212,18 +43342,22 @@ }, "nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.1e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, "output_cost_per_token_flex": 1.2e-07, + "output_cost_per_token_priority": 4.2e-07, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_native_structured_output": true, "supports_response_schema": true, @@ -43231,15 +43365,19 @@ }, "nvidia.nemotron-super-3-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.5e-07, + "output_cost_per_token_batches": 3.25e-07, "output_cost_per_token_flex": 3.25e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.1375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_reasoning": true, @@ -44731,6 +44869,7 @@ "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", @@ -44739,9 +44878,10 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, "output_cost_per_token_priority": 1.05e-06, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -44754,6 +44894,7 @@ "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, "input_cost_per_token_priority": 1.225e-07, "litellm_provider": "bedrock_converse", @@ -44762,9 +44903,10 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3e-07, + "output_cost_per_token_batches": 1.5e-07, "output_cost_per_token_flex": 1.5e-07, "output_cost_per_token_priority": 5.25e-07, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -44772,15 +44914,19 @@ }, "openai.gpt-oss-safeguard-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, @@ -44790,15 +44936,19 @@ }, "openai.gpt-oss-safeguard-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.2e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_batches": 1e-07, "output_cost_per_token_flex": 1e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_bedrock_runtime_chat_completions_inline_reasoning": true, "supports_function_calling": true, @@ -45654,13 +45804,17 @@ "qwen.qwen3-coder-30b-a3b-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 262144, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -45668,15 +45822,19 @@ }, "qwen.qwen3-32b-v1:0": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 32768, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -45796,13 +45954,17 @@ }, "qwen.qwen3-vl-235b-a22b": { "input_cost_per_token": 5.3e-07, + "input_cost_per_token_batches": 2.6e-07, + "input_cost_per_token_priority": 9.3e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.66e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_batches": 1.33e-06, + "output_cost_per_token_priority": 4.66e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_function_calling": true, "supports_system_messages": true, @@ -45812,18 +45974,22 @@ }, "qwen.qwen3-coder-next": { "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 8.75e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 256000, "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", "output_cost_per_token": 1.2e-06, + "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -47842,12 +48008,15 @@ "us.amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -47858,12 +48027,15 @@ "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -47872,15 +48044,21 @@ "us.amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -48599,14 +48777,19 @@ }, "us.deepseek.v3.2": { "input_cost_per_token": 6.2e-07, + "input_cost_per_token_batches": 3.1e-07, "input_cost_per_token_flex": 3.1e-07, + "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 1.85e-06, + "output_cost_per_token_batches": 9.25e-07, "output_cost_per_token_flex": 9.25e-07, + "output_cost_per_token_priority": 3.2375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, @@ -48648,26 +48831,30 @@ "us.meta.llama3-1-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-1-8b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-1-8b-instruct.md", + "output_cost_per_token_batches": 1.1e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -48722,13 +48909,15 @@ "us.meta.llama3-3-70b-instruct-v1:0": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7.2e-07, + "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.2e-07, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-3-70b-instruct.md", + "output_cost_per_token_batches": 3.6e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_tool_choice": false }, @@ -53430,37 +53619,45 @@ "zai.glm-4.7": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 203000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.2e-06, + "output_cost_per_token_batches": 1.1e-06, "output_cost_per_token_flex": 1.1e-06, + "output_cost_per_token_priority": 3.85e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-zai-glm-4-7.md", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false }, "zai.glm-5": { "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_priority": 1.75e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, "supports_function_calling": true, "supports_reasoning": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -53468,19 +53665,23 @@ "zai.glm-4.7-flash": { "deprecation_date": "2027-03-30", "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, + "input_cost_per_token_priority": 1.225e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 203000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "output_cost_per_token_priority": 7e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-zai-glm-4-7-flash.md", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false @@ -58851,7 +59052,9 @@ }, "bedrock_mantle/openai.gpt-oss-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -58859,7 +59062,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-120b.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -58872,7 +59077,9 @@ }, "bedrock_mantle/openai.gpt-oss-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3.5e-08, "input_cost_per_token_flex": 3.5e-08, + "input_cost_per_token_priority": 1.225e-07, "output_cost_per_token": 3e-07, "output_cost_per_token_flex": 1.5e-07, "litellm_provider": "bedrock_mantle", @@ -58880,7 +59087,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-20b.html", + "output_cost_per_token_batches": 1.5e-07, + "output_cost_per_token_priority": 5.25e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -58893,7 +59102,9 @@ }, "bedrock_mantle/openai.gpt-oss-safeguard-120b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -58901,7 +59112,9 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-safeguard-120b.html", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -58912,7 +59125,9 @@ }, "bedrock_mantle/openai.gpt-oss-safeguard-20b": { "input_cost_per_token": 7e-08, + "input_cost_per_token_batches": 3e-08, "input_cost_per_token_flex": 3e-08, + "input_cost_per_token_priority": 1.2e-07, "output_cost_per_token": 2e-07, "output_cost_per_token_flex": 1e-07, "litellm_provider": "bedrock_mantle", @@ -58920,7 +59135,9 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-oss-safeguard-20b.html", + "output_cost_per_token_batches": 1e-07, + "output_cost_per_token_priority": 3.5e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supported_endpoints": [ "/v1/chat/completions" ], @@ -59553,7 +59770,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_prompt_caching": true, "supports_prompt_cache_breakpoint": true, "supports_reasoning": true, @@ -59592,7 +59809,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_prompt_caching": true, "supports_prompt_cache_breakpoint": true, "supports_reasoning": true, @@ -59675,7 +59892,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59714,7 +59931,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59799,7 +60016,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59834,7 +60051,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59872,7 +60089,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -59907,7 +60124,7 @@ "supports_function_calling": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, - "supports_none_reasoning_effort": false, + "supports_none_reasoning_effort": true, "supports_tool_choice": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -60002,7 +60219,9 @@ }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.45e-07, "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, "litellm_provider": "bedrock_mantle", @@ -60010,6 +60229,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_priority": 7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60023,7 +60245,9 @@ }, "bedrock_mantle/google.gemma-4-26b-a4b": { "input_cost_per_token": 1.3e-07, + "input_cost_per_token_batches": 6.5e-08, "input_cost_per_token_flex": 6.5e-08, + "input_cost_per_token_priority": 2.275e-07, "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, "litellm_provider": "bedrock_mantle", @@ -60031,6 +60255,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", + "output_cost_per_token_batches": 2e-07, + "output_cost_per_token_priority": 7e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60044,7 +60271,9 @@ }, "bedrock_mantle/google.gemma-4-e2b": { "input_cost_per_token": 4e-08, + "input_cost_per_token_batches": 2e-08, "input_cost_per_token_flex": 2e-08, + "input_cost_per_token_priority": 7e-08, "output_cost_per_token": 8e-08, "output_cost_per_token_flex": 4e-08, "litellm_provider": "bedrock_mantle", @@ -60052,6 +60281,9 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_token_batches": 4e-08, + "output_cost_per_token_priority": 1.4e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60068,8 +60300,10 @@ "input_cost_per_token": 1.25e-06, "output_cost_per_token": 2.5e-06, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_mantle", @@ -60077,6 +60311,7 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", + "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.375e-06, "supported_endpoints": [ @@ -60088,7 +60323,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/xai.grok-4.6": { "use_openai_responses_path": true, @@ -67524,7 +67759,9 @@ "bedrock_mantle/deepseek.v3.1": { "deprecation_date": "2027-03-30", "input_cost_per_token": 5.8e-07, + "input_cost_per_token_batches": 2.9e-07, "input_cost_per_token_flex": 2.9e-07, + "input_cost_per_token_priority": 1.015e-06, "output_cost_per_token": 1.68e-06, "output_cost_per_token_flex": 8.4e-07, "litellm_provider": "bedrock_mantle", @@ -67532,17 +67769,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 8.4e-07, + "output_cost_per_token_priority": 2.94e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-1.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/moonshotai.kimi-k2-thinking": { "deprecation_date": "2027-03-30", "input_cost_per_token": 6e-07, + "input_cost_per_token_batches": 3e-07, "input_cost_per_token_flex": 3e-07, + "input_cost_per_token_priority": 1.05e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_flex": 1.25e-06, "litellm_provider": "bedrock_mantle", @@ -67550,17 +67791,21 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.38e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-235b-a22b-2507": { "input_cost_per_token": 2.2e-07, + "input_cost_per_token_batches": 1.1e-07, "input_cost_per_token_flex": 1.1e-07, + "input_cost_per_token_priority": 3.85e-07, "output_cost_per_token": 8.8e-07, "output_cost_per_token_flex": 4.4e-07, "litellm_provider": "bedrock_mantle", @@ -67568,17 +67813,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 4.4e-07, + "output_cost_per_token_priority": 1.54e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-235b-a22b-2507.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-32b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -67586,18 +67835,22 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-32b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-coder-30b-a3b-instruct": { "deprecation_date": "2027-03-30", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "output_cost_per_token": 6e-07, "output_cost_per_token_flex": 3e-07, "litellm_provider": "bedrock_mantle", @@ -67605,17 +67858,21 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 3e-07, + "output_cost_per_token_priority": 1.05e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-30b-a3b-instruct.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-coder-480b-a35b-instruct": { "deprecation_date": "2027-03-30", "input_cost_per_token": 4.5e-07, + "input_cost_per_token_batches": 2.25e-07, "input_cost_per_token_flex": 2.25e-07, + "input_cost_per_token_priority": 7.875e-07, "output_cost_per_token": 1.8e-06, "output_cost_per_token_flex": 9e-07, "litellm_provider": "bedrock_mantle", @@ -67623,16 +67880,20 @@ "max_output_tokens": 16000, "max_tokens": 16000, "mode": "chat", + "output_cost_per_token_batches": 9e-07, + "output_cost_per_token_priority": 3.15e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-coder-480b-a35b-instruct.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-next-80b-a3b-instruct": { "input_cost_per_token": 1.4e-07, + "input_cost_per_token_batches": 7e-08, "input_cost_per_token_flex": 7e-08, + "input_cost_per_token_priority": 2.6e-07, "output_cost_per_token": 1.2e-06, "output_cost_per_token_flex": 6e-07, "litellm_provider": "bedrock_mantle", @@ -67640,17 +67901,21 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_priority": 2.1e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_reasoning": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-next-80b-a3b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "bedrock_mantle/qwen.qwen3-vl-235b-a22b-instruct": { "input_cost_per_token": 5.3e-07, + "input_cost_per_token_batches": 2.6e-07, "input_cost_per_token_flex": 2.6e-07, + "input_cost_per_token_priority": 9.3e-07, "output_cost_per_token": 2.66e-06, "output_cost_per_token_flex": 1.33e-06, "litellm_provider": "bedrock_mantle", @@ -67658,13 +67923,15 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", + "output_cost_per_token_batches": 1.33e-06, + "output_cost_per_token_priority": 4.66e-06, "supported_endpoints": [ "/v1/chat/completions" ], "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-qwen-qwen3-vl-235b-a22b.html" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json" }, "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, @@ -75218,5 +75485,45 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "baseten/reflection/Beam-1.0": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 1.3e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.3e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "nebius/reflection/beam": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/reflection%2Fbeam", + "supports_function_calling": true, + "supports_reasoning": true } } diff --git a/pyproject.toml b/pyproject.toml index 0be8adb7337..854811799a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -290,7 +290,7 @@ ci = [ "langchain-openai==1.1.14", "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", - "claude-agent-sdk==0.1.44", + "claude-agent-sdk==0.2.165", "google-cloud-speech==2.40.0", ] healthcheck = [ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index e4999c71506..681476456c2 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -112,6 +112,8 @@ ignored_function_names = [ "_get_wildcard_deployments", # Tested through the get_model_list_of_routed_group wildcard test in test_router.py "_is_fallback_hop", # Tested through the order fallback hop tests in test_router_order_fallback.py "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py + "_merged_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py + "_drops_malformed_request_tools", # Tested through the _update_kwargs_with_deployment tools merge tests in test_router.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_register_continuation_prefill_check", # Tested through the continuation tests in test_router.py diff --git a/tests/e2e/llm_translation/responses_helpers.py b/tests/e2e/llm_translation/responses_helpers.py index be0fac74bb6..29b923a7f34 100644 --- a/tests/e2e/llm_translation/responses_helpers.py +++ b/tests/e2e/llm_translation/responses_helpers.py @@ -8,10 +8,13 @@ AZURE_OPENAI_BACKEND: Final = "azure/gpt-5.4-nano" AZURE_OPENAI_API_VERSION: Final = "v1" -def azure_openai_params(api_version: str = AZURE_OPENAI_API_VERSION) -> LiteLLMParamsBody: +def azure_openai_params( + api_version: str = AZURE_OPENAI_API_VERSION, responses_api_version: str | None = None +) -> LiteLLMParamsBody: return LiteLLMParamsBody( model=AZURE_OPENAI_BACKEND, api_base="os.environ/AZURE_API_BASE", api_key="os.environ/AZURE_API_KEY", api_version=api_version, + responses_api_version=responses_api_version, ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 15cf20aa65a..46561ccbc37 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -24,7 +24,10 @@ from sdk_clients import SdkClients pytestmark = pytest.mark.e2e OPENAI_IMAGE_BACKEND: Final = "openai/gpt-image-1-mini" -BEDROCK_IMAGE_BACKEND: Final = "bedrock/amazon.nova-canvas-v1:0" +BEDROCK_IMAGE_BACKEND: Final = "bedrock/stability.stable-image-core-v1:1" +# Nova Canvas reached end of life on 2026-09-30. Stable Image Core is active on Bedrock only in us-west-2: +# https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-stability-ai-stable-image-core-1-0.html (2026-10-16) +BEDROCK_IMAGE_REGION: Final = "us-west-2" class _OptionalImageBody(BaseModel): @@ -96,10 +99,10 @@ class TestImageGeneration: model=BEDROCK_IMAGE_BACKEND, aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", - aws_region_name="os.environ/AWS_REGION", + aws_region_name=BEDROCK_IMAGE_REGION, ), ) - images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat") _assert_image_returned(images) @pytest.mark.skip(reason="stage red: product gap, /v1/images/generations 500s (aimage_generation TypeError) on missing prompt instead of 400") diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index ee4699703b2..bb75569ad99 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -40,7 +40,7 @@ from provider_edge import LiveEdge, start_provider_edge from provider_edge_bedrock import ConverseBodyCapture, bedrock_signer from proxy_client import ProxyClient from pydantic import BaseModel, TypeAdapter -from responses_helpers import AZURE_OPENAI_BACKEND, azure_openai_params +from responses_helpers import AZURE_OPENAI_API_VERSION, AZURE_OPENAI_BACKEND, azure_openai_params from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -748,6 +748,70 @@ class TestResponses: f"/responses over azure openai preview returned no output text: {response.output!r}" ) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.AZURE,), + models=(AZURE_OPENAI_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_azure_openai_dated_ga_chat_api_version_serves_responses( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register( + proxy, + resources, + azure_openai_params(api_version="2024-10-21"), + prefix="e2e-responses-azure-dated-chat", + ) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + extra_body=NO_PROXY_CACHE, + ) + assert response.output_text.strip(), ( + f"/responses over an azure deployment pinned to a dated GA chat api_version returned no output text: " + f"{response.output!r}" + ) + + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.AZURE,), + models=(AZURE_OPENAI_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_azure_openai_responses_api_version_overrides_a_pre_responses_api_version( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + # Azure rejects /openai/responses below api-version 2025-03-01-preview (probed live 2026-10-15), + # so this call only succeeds if the proxy sends responses_api_version instead of api_version. + model: Final = _register( + proxy, + resources, + azure_openai_params(api_version="2024-10-01-preview", responses_api_version=AZURE_OPENAI_API_VERSION), + prefix="e2e-responses-azure-override", + ) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + extra_body=NO_PROXY_CACHE, + ) + assert response.output_text.strip(), ( + f"/responses over an azure deployment with responses_api_version set returned no output text: " + f"{response.output!r}" + ) + @meta( Subject( domain=Domain.LLM_TRANSLATION, diff --git a/tests/e2e/mcp/test_mcp_tool_search_e2e.py b/tests/e2e/mcp/test_mcp_tool_search_e2e.py index 847239e8a66..c4a8e29eb82 100644 --- a/tests/e2e/mcp/test_mcp_tool_search_e2e.py +++ b/tests/e2e/mcp/test_mcp_tool_search_e2e.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Final import pytest +from pydantic import BaseModel, ConfigDict, TypeAdapter from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import provider_edge_base, unique_marker from e2e_http import unwrap @@ -17,6 +18,19 @@ MCP_TOOL_SEARCH_TOOL_NAME: Final = "mcp_tool_search" OVERSIZED_DESCRIPTION: Final = f"OVERSIZED {'weather ' * 12_000}" +class _SearchHit(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + + +_SEARCH_HITS: Final = TypeAdapter(tuple[_SearchHit, ...]) + + +def _hits_tool(hits: tuple[_SearchHit, ...], tool_name: str) -> bool: + return any(hit.name.endswith(f"-{tool_name}") for hit in hits) + + class TestMcpToolSearch: @pytest.mark.covers("mcp.tool_search.api_key.oversized_description_omitted") def test_native_tool_search_omits_oversized_description( @@ -82,5 +96,6 @@ class TestMcpToolSearch: ) assert result.is_error is not True, result.all_text - assert normal_tool_name in result.all_text - assert oversized_tool_name not in result.all_text + hits: Final = _SEARCH_HITS.validate_json(result.all_text) + assert _hits_tool(hits, normal_tool_name), hits + assert not _hits_tool(hits, oversized_tool_name), hits diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 08fc7a79df3..d7a10c6da1f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1366,6 +1366,7 @@ class LiteLLMParamsBody(BaseModel): litellm_credential_name: str | None = None api_base: str | None = Field(default=None, repr=False) api_version: str | None = None + responses_api_version: str | None = None realtime_protocol: str | None = None allowed_openai_params: list[str] | None = None drop_params: bool | None = None diff --git a/tests/e2e/ui/helpers/playground.ts b/tests/e2e/ui/helpers/playground.ts index 866334e1c94..b46a2263072 100644 --- a/tests/e2e/ui/helpers/playground.ts +++ b/tests/e2e/ui/helpers/playground.ts @@ -2,42 +2,57 @@ import { expect, type Locator, type Page as PlaywrightPage } from "@playwright/t import { navigateToPage, dismissFeedbackPopup } from "./navigation"; import { Page } from "../fixtures/pages"; -/** Controls for the Test Key / Playground page, shared with the router-fallback specs. */ - -/** - * The configuration panel is rendered twice, docked and overlay, with one visible at a time. - * Every control is narrowed to the visible copy or it trips strict mode against its hidden twin. - */ export const onlyVisible = (locator: Locator): Locator => locator.filter({ visible: true }).first(); -/** The model combobox, addressed by the placeholder its search input shows before selection. */ -export const modelSelect = (page: PlaywrightPage): Locator => onlyVisible(page.getByPlaceholder("Select a Model")); +export const composer = (page: PlaywrightPage): Locator => + onlyVisible(page.getByRole("group").filter({ has: page.getByRole("textbox", { name: "Message" }) })); + +export const modelSelect = (page: PlaywrightPage): Locator => + onlyVisible(page.getByRole("combobox", { name: "Model", exact: true })); export const sendButton = (page: PlaywrightPage): Locator => - onlyVisible(page.getByRole("button", { name: "Send message" })); + composer(page).getByRole("button", { name: "Send message" }); -/** The Virtual Key Source dropdown, addressed by the accessible name on its trigger. */ -export const keySourceSelect = (page: PlaywrightPage): Locator => - onlyVisible(page.getByLabel("Virtual Key Source")); +export const keySourceSelect = (page: PlaywrightPage): Locator => onlyVisible(page.getByLabel("Send as")); export async function openPlayground(page: PlaywrightPage): Promise { await navigateToPage(page, Page.LlmPlayground); await dismissFeedbackPopup(page); - await expect(onlyVisible(page.getByText("Virtual Key Source"))).toBeVisible({ - timeout: 20_000, - }); + await expect(modelSelect(page)).toBeVisible({ timeout: 20_000 }); + const searchHint = page.getByRole("button", { name: "Dismiss search hint" }); + if (await searchHint.isVisible()) { + await searchHint.click(); + } +} + +export async function openRequestSettings(page: PlaywrightPage): Promise { + await composer(page).getByRole("button", { name: "Settings", exact: true }).click(); + await expect(keySourceSelect(page)).toBeVisible({ timeout: 10_000 }); +} + +export async function enterVirtualKey(page: PlaywrightPage, key: string): Promise { + await openRequestSettings(page); + await keySourceSelect(page).click(); + await onlyVisible(page.getByRole("option", { name: "Virtual key", exact: true })).click({ timeout: 15_000 }); + const keyInput = onlyVisible(page.getByLabel("Virtual key", { exact: true })); + await expect(keyInput).toBeVisible({ timeout: 10_000 }); + await keyInput.fill(key); + await page.keyboard.press("Escape"); +} + +export async function searchModels(page: PlaywrightPage, query: string): Promise { + await modelSelect(page).click(); + await onlyVisible(page.getByPlaceholder("Search models")).fill(query); } export async function selectModel(page: PlaywrightPage, model: string): Promise { - const select = modelSelect(page); - await select.click(); - // Virtualized: options outside the rendered window are absent from the DOM, so search first. - await select.fill(model); + await searchModels(page, model); await onlyVisible(page.getByRole("option", { name: model, exact: true })).click({ timeout: 15_000 }); + await expect(modelSelect(page)).toContainText(model); } export async function sendMessage(page: PlaywrightPage, message: string): Promise { - const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + const input = composer(page).getByRole("textbox", { name: "Message" }); await expect(input).toBeVisible({ timeout: 15_000 }); await input.fill(message); await sendButton(page).click(); diff --git a/tests/e2e/ui/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts index fcdd71898e2..8d251d52ef7 100644 --- a/tests/e2e/ui/tests/auth/logout.spec.ts +++ b/tests/e2e/ui/tests/auth/logout.spec.ts @@ -2,6 +2,8 @@ import { test, expect } from "@playwright/test"; import { Role, users } from "../../fixtures/users"; import { logInThroughLoginPage } from "../../helpers/userOnboarding"; +const LOGOUT_URL = process.env.PROXY_LOGOUT_URL ?? ""; + test.describe("Logout", () => { test.use({ storageState: { cookies: [], origins: [] } }); @@ -9,9 +11,21 @@ test.describe("Logout", () => { const admin = users[Role.ProxyAdmin]; await logInThroughLoginPage(page, admin.email, admin.password); + if (LOGOUT_URL) { + const target = new URL(LOGOUT_URL); + await page.route( + (url) => url.origin === target.origin, + (route) => + route.fulfill({ status: 200, contentType: "text/html", body: "logged out" }), + ); + } + const settingsLoaded = page.waitForResponse((r) => r.url().includes("/sso/get/ui_settings") && r.ok(), { + timeout: 30_000, + }); await page.goto("/ui"); // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); + await settingsLoaded; // Open the sidebar account menu. The trigger button exposes an aria-label // of "Account menu — — signed in as "; clicking it opens the @@ -22,8 +36,12 @@ test.describe("Logout", () => { await expect(popup).toBeVisible({ timeout: 5_000 }); // Click Logout — the handler clears the auth cookie and navigates via - // window.location.href = PROXY_LOGOUT_URL (empty string in the e2e env). + // window.location.href = PROXY_LOGOUT_URL. + const loggedOut = LOGOUT_URL + ? page.waitForURL((url) => url.origin === new URL(LOGOUT_URL).origin, { timeout: 15_000 }) + : page.waitForEvent("load", { timeout: 15_000 }); await popup.getByRole("button", { name: "Logout" }).click(); + await loggedOut; await expect .poll(async () => (await page.context().cookies()).filter((c) => c.name === "token").length, { timeout: 15_000, diff --git a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts index b4dc7c4e8da..c57a9147ab0 100644 --- a/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts +++ b/tests/e2e/ui/tests/guardrails/presidioUserStory.spec.ts @@ -3,7 +3,7 @@ import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL, PROPAGATION_TIMEOUT_MS } from ". import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; import { CHAT_MODEL_A, masterKey, rootPath, uniqueSuffix, waitForSpendLog } from "../../helpers/traffic"; -import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; +import { composer, openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground"; const RAW_EMAIL = "jane.doe@example.com"; const RAW_PHONE = "555-867-5309"; @@ -123,14 +123,12 @@ test.describe("Presidio PII guardrail, end to end from the dashboard", () => { await openPlayground(page); await selectModel(page, CHAT_MODEL_A); - const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails")); - await expect(guardrailSelect).toBeVisible({ timeout: 20_000 }); - await guardrailSelect.click(); - await guardrailSelect.fill(guardrailName); + await composer(page).getByRole("combobox", { name: "Guardrails" }).click(); + await onlyVisible(page.getByPlaceholder("Search guardrails")).fill(guardrailName); await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 }); await page.keyboard.press("Escape"); - const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false })); + const input = composer(page).getByRole("textbox", { name: "Message" }); await expect(input).toBeVisible({ timeout: 15_000 }); await input.fill(prompt); diff --git a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts index f392c5104da..d34bc034453 100644 --- a/tests/e2e/ui/tests/internal-user/internalUser.spec.ts +++ b/tests/e2e/ui/tests/internal-user/internalUser.spec.ts @@ -9,7 +9,7 @@ import { import { Page } from "../../fixtures/pages"; import { navigateToPage, clickTeamId } from "../../helpers/navigation"; import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic"; -import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; +import { openPlayground, selectModel, sendMessage, enterVirtualKey } from "../../helpers/playground"; test.describe("Internal User", () => { test.use({ storageState: INTERNAL_USER_STORAGE_PATH }); @@ -73,12 +73,7 @@ test.describe("Internal User", () => { try { await openPlayground(page); - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); - - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(apiKey); + await enterVirtualKey(page, apiKey); await selectModel(page, CHAT_MODEL_A); await sendMessage(page, `internal user team key ping ${keyName}`); diff --git a/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts b/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts new file mode 100644 index 00000000000..7d0f02209cf --- /dev/null +++ b/tests/e2e/ui/tests/modelsPage/cacheWriteTiers.spec.ts @@ -0,0 +1,146 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { captureRequestBody, readBack } from "../../helpers/roundTrip"; + +const FIVE_MINUTE_PER_MILLION = 3.75; +const ONE_HOUR_PER_MILLION = 7.5; +const perToken = (perMillion: number): number => perMillion / 1_000_000; +const ONE_HOUR_LABEL = /^Cache Write Cost \(1 hour, per 1M tokens\)/; + +type Deployment = { + model_name?: string; + model_info?: { id?: string }; + litellm_params: { + cache_creation_input_token_cost?: number; + cache_creation_input_token_cost_above_1hr?: number; + }; +}; + +const masterKey = users[Role.ProxyAdmin].password; +const auth = { Authorization: `Bearer ${masterKey}` }; + +async function deleteModel(page: PlaywrightPage, id: string) { + await page.request.post("/model/delete", { headers: auth, data: { id } }); +} + +async function readDeployment(page: PlaywrightPage, modelId: string): Promise { + const body = await readBack<{ data: Deployment[] }>(page, `/model/info?litellm_model_id=${modelId}`); + return body.data[0]; +} + +async function selectProvider(page: PlaywrightPage, providerName: string) { + const providerDropdown = page.getByRole("combobox", { name: "Provider", exact: true }); + await providerDropdown.click(); + await providerDropdown.fill(providerName); + await page.getByRole("option").filter({ hasText: providerName }).first().click(); + await expect(providerDropdown).toHaveValue(providerName); +} + +test.describe("Cache write cost tiers", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + let createdModelId: string | null = null; + + test.afterEach(async ({ page }) => { + if (createdModelId) { + await deleteModel(page, createdModelId); + createdModelId = null; + } + }); + + test("the edit form saves the 1 hour rate next to the 5 minute rate and the detail view shows both", async ({ + page, + }) => { + const modelName = `e2e-cache-tiers-edit-${Date.now()}`; + const created = await page.request.post("/model/new", { + headers: auth, + data: { + model_name: modelName, + litellm_params: { + model: "anthropic/claude-sonnet-4-6", + api_key: "sk-e2e-not-used", + cache_creation_input_token_cost: perToken(FIVE_MINUTE_PER_MILLION), + }, + model_info: {}, + }, + }); + expect(created.ok(), `POST /model/new for ${modelName}`).toBe(true); + createdModelId = (await created.json()).model_info?.id; + expect(createdModelId, "model id from /model/new").toBeTruthy(); + + await page.goto("/ui"); + await page.getByText("Models + Endpoints").click(); + const modelIdCell = page.getByTestId(`model-id-${createdModelId}`); + await expect(modelIdCell).toBeVisible({ timeout: 15_000 }); + await modelIdCell.click(); + await expect(page.getByText("Back to Models").first()).toBeVisible({ timeout: 10_000 }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + const oneHour = page.getByLabel(ONE_HOUR_LABEL); + await oneHour.waitFor({ timeout: 15_000 }); + await oneHour.fill(String(ONE_HOUR_PER_MILLION)); + + const patchBody = await captureRequestBody( + page, + { method: "PATCH", urlIncludes: `/model/${createdModelId}/update` }, + async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }, + ); + expect(patchBody.litellm_params.cache_creation_input_token_cost_above_1hr, "1 hour rate sent per token").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + await expect(page.getByText("Model settings updated successfully")).toBeVisible({ timeout: 10_000 }); + + const stored = await readDeployment(page, createdModelId!); + expect(stored?.litellm_params?.cache_creation_input_token_cost_above_1hr, "1 hour rate persisted").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + expect(stored?.litellm_params?.cache_creation_input_token_cost, "5 minute rate untouched").toBe( + perToken(FIVE_MINUTE_PER_MILLION), + ); + + await expect(page.getByText(ONE_HOUR_LABEL)).toBeVisible(); + await expect(page.getByText(ONE_HOUR_PER_MILLION.toFixed(4))).toBeVisible({ timeout: 10_000 }); + await expect(page.getByText(FIVE_MINUTE_PER_MILLION.toFixed(4))).toBeVisible(); + }); + + test("Add Model sends both cache write tiers per token", async ({ page }) => { + const publicName = `e2e-cache-tiers-add-${Date.now()}`; + await navigateToPage(page, Page.Models); + await page.getByRole("tab", { name: "Add Model" }).click(); + await selectProvider(page, "Anthropic"); + + await page.getByRole("combobox", { name: "Select models" }).click(); + await page.getByRole("option", { name: "Custom Model Name (Enter below)" }).click(); + await page.keyboard.press("Escape"); + await page.getByPlaceholder("Enter custom model name").fill(publicName); + await page.locator("#api_key").fill("sk-e2e-not-used"); + + await page.getByText("Advanced Settings").click(); + await page.getByRole("switch", { name: "Custom Pricing" }).click(); + await page.getByLabel(/^Input Cost \(per 1M tokens\)/).fill("3"); + await page.getByLabel(/^Output Cost \(per 1M tokens\)/).fill("15"); + await page.getByLabel(/^Cache Write Cost \(per 1M tokens\)/).fill(String(FIVE_MINUTE_PER_MILLION)); + await page.getByLabel(ONE_HOUR_LABEL).fill(String(ONE_HOUR_PER_MILLION)); + + const body = await captureRequestBody(page, { method: "POST", urlIncludes: "/model/new" }, async () => { + await page.getByTestId("add-model-btn").click(); + }); + await expect(page.getByText("created successfully")).toBeVisible({ timeout: 15_000 }); + const listed = await readBack<{ data: Deployment[] }>(page, "/v2/model/info"); + createdModelId = listed.data.find((row) => row.model_name === publicName)?.model_info?.id ?? null; + + expect(body.model_name).toBe(publicName); + expect(body.litellm_params.cache_creation_input_token_cost, "5 minute rate sent per token").toBe( + perToken(FIVE_MINUTE_PER_MILLION), + ); + expect(body.litellm_params.cache_creation_input_token_cost_above_1hr, "1 hour rate sent per token").toBe( + perToken(ONE_HOUR_PER_MILLION), + ); + expect(createdModelId, "the model landed in the database").toBeTruthy(); + }); +}); diff --git a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts index c532641b238..34a93c7ce14 100644 --- a/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts +++ b/tests/e2e/ui/tests/modelsPage/clearCustomPricing.spec.ts @@ -87,7 +87,7 @@ test.describe("Clear custom pricing on a deployment", () => { // Both cache fields share the same placeholder ("Defaults to Input Cost if blank"), // so disambiguate via their labels. const cacheReadCost = page.getByLabel(/Cache Read Cost/); - const cacheWriteCost = page.getByLabel(/Cache Write Cost/); + const cacheWriteCost = page.getByLabel("Cache Write Cost (per 1M tokens)", { exact: true }); await inputCost.waitFor({ timeout: 15_000 }); for (const field of [inputCost, outputCost, cacheReadCost, cacheWriteCost]) { await field.click({ clickCount: 3 }); diff --git a/tests/e2e/ui/tests/playground/playground.spec.ts b/tests/e2e/ui/tests/playground/playground.spec.ts index 9a9ecd539b3..205cc57ae46 100644 --- a/tests/e2e/ui/tests/playground/playground.spec.ts +++ b/tests/e2e/ui/tests/playground/playground.spec.ts @@ -1,7 +1,14 @@ import { test, expect } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; import { CHAT_MODEL_A, CHAT_MODEL_B, MOCK_RESPONSE_TEXT, createVirtualKey } from "../../helpers/traffic"; -import { keySourceSelect, onlyVisible, openPlayground, selectModel, sendMessage } from "../../helpers/playground"; +import { + keySourceSelect, + openPlayground, + openRequestSettings, + selectModel, + sendMessage, + enterVirtualKey, +} from "../../helpers/playground"; /** * The one flow that exercises the dashboard's own LLM call path rather than an admin CRUD endpoint, @@ -14,8 +21,9 @@ test.describe("Playground", () => { test(`chats with ${model} using the current UI session`, async ({ page }) => { await openPlayground(page); - // "Current UI Session" is the default: the logged-in admin's key, nothing pasted. - await expect(keySourceSelect(page)).toContainText("Current UI Session"); + await openRequestSettings(page); + await expect(keySourceSelect(page)).toContainText("Current UI session"); + await page.keyboard.press("Escape"); await selectModel(page, model); const prompt = `playground ping for ${model}`; @@ -34,14 +42,7 @@ test.describe("Playground", () => { await openPlayground(page); - // Switch the source to "Virtual Key" and paste the key we just minted. - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); - - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(key); - + await enterVirtualKey(page, key); await selectModel(page, CHAT_MODEL_A); await sendMessage(page, "playground ping via virtual key"); diff --git a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts index f3c031f0172..ff41778b2c9 100644 --- a/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts +++ b/tests/e2e/ui/tests/team-admin/teamAdmin.spec.ts @@ -10,7 +10,7 @@ import { Page } from "../../fixtures/pages"; import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation"; import { captureRequestBody, readBack } from "../../helpers/roundTrip"; import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic"; -import { keySourceSelect, modelSelect, onlyVisible, openPlayground } from "../../helpers/playground"; +import { onlyVisible, openPlayground, searchModels, enterVirtualKey } from "../../helpers/playground"; /** * Every identifier a roster is addressable by. Which of user_id / user_email is populated depends on @@ -228,21 +228,14 @@ test.describe("Team Admin", () => { .toBe(true); await openPlayground(page); - await keySourceSelect(page).click(); - await onlyVisible(page.getByRole("option", { name: "Virtual Key" })).click({ timeout: 15_000 }); + await enterVirtualKey(page, teamKey); - const keyInput = onlyVisible(page.getByPlaceholder("Enter custom Virtual Key")); - await expect(keyInput).toBeVisible({ timeout: 10_000 }); - await keyInput.fill(teamKey); - - const select = modelSelect(page); - await select.click(); - await select.fill(teamModelName); + await searchModels(page, teamModelName); await expect(onlyVisible(page.getByRole("option", { name: teamModelName }))).toBeVisible({ timeout: 15_000, }); - await select.fill(CHAT_MODEL_A); + await onlyVisible(page.getByPlaceholder("Search models")).fill(CHAT_MODEL_A); await expect(onlyVisible(page.getByRole("option", { name: CHAT_MODEL_A }))).toBeVisible({ timeout: 15_000, }); diff --git a/tests/integration/AGENTS.md b/tests/integration/AGENTS.md index 5a09b66e7ed..a727b20596c 100644 --- a/tests/integration/AGENTS.md +++ b/tests/integration/AGENTS.md @@ -25,3 +25,15 @@ By the domain a user would name: `pricing`, `spend`, `routing`, `mcp`. A file on directory that a `GROUPS` entry in `run.py` selects. There is no manifest and no `covers` marker on new tests. A product bug the test exposes is `pytest.skip("BUG: ")` at the top of the body, not a fix in the test and not a deletion. Needs no network, proxy, DB or Redis: `tests/unit` + +CircleCI's integration groups run every file here except the OCI and sandbox tests, which no CI runs, and +the files in `GITHUB_FILES` in `run.py`. CircleCI runs on main, not on pull requests, so a new file here +first runs after its PR merges. Wanting a file to run before merge is no reason to change CI, and a new +dedicated workflow for it is the same mistake as a shared shard. Write the file, leave `GITHUB_FILES` and +every workflow alone, and tell the user it first runs on main. `GITHUB_FILES` is only for a file that +needs what CircleCI's groups cannot give it, like the images `image-scan.yml` builds or the plain pytest +run CircleCI's `roi-database` job gives `test_roi_observed.py`. Add a file there only in that case, +together with its dedicated job, and list a new dedicated GitHub workflow with the directory it owns in +`DEDICATED_INTEGRATION_WORKFLOWS` in `.github/scripts/assert_ci_coverage.py`. The `ci-coverage` check +fails a file here that a GitHub workflow runs unless it is in `GITHUB_FILES` and that workflow owns its +directory (#46780) diff --git a/tests/integration/README.md b/tests/integration/README.md index 33263043489..496c88050ee 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -4,7 +4,7 @@ These tests exercise a running gateway, PostgreSQL and Redis with an owned local The ROI database contracts in `database/test_roi_observed.py` run as plain pytest outside `run_integration.sh` in CircleCI's `roi-database` Postgres job. The job uploads coverage with the `roi-postgres` flag. These contracts own temporary databases and script only the external provider transport -The Redis cluster IAM contracts in `sdk/test_redis_cluster_iam_auth.py` start their own `redis-server` and mock the cloud credential clients, so they run in the GitHub Actions `redis-compat` matrix against every supported redis-py major rather than in `run_integration.sh` +The Redis cluster IAM contracts in `sdk/test_redis_cluster_iam_auth.py` start their own `redis-server` and mock the cloud credential clients, so CircleCI's `sdk` group installs `redis-server` before it runs them The `cost` group is driven by the `CostTrackingTestCase` literals under `cost_calculation/cases/.py` (gathered by `cost_calculation/catalog.py`), with the scripted cost map in `cost_calculation/cost_map.py`. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, and add hand-computed expected values. The upstream serves each stored response for any path under `/`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL` @@ -26,7 +26,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r Owned proxies disable `.env` loading by default. A contract testing dotenv behavior can explicitly set `PYTHON_DOTENV_DISABLED=0` in its owned environment. The runner drops inherited `*_API_KEY` variables, so a contract that needs one sets a synthetic value in its owned environment. An explicit scratch writer database drops an inherited reader with a different database name; pass both database URLs when the contract needs a paired reader. Database relays close listeners, active connections and their event loops on teardown. Upstream startup failures reap the child process before returning the failure. Custom subprocesses, direct SQL and background tasks still need an explicit owner in the test -Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended +Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. The `kill()` that `replacing_worker` yields returns as soon as the worker is dead, so a test can release requests it holds on the surviving worker right away. Uvicorn may take a full worker healthcheck to notice the death, so the block waits for the supervisor to reap the worker and for the replacement only when it exits. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended Any port the harness picks before the server that listens on it has started comes from `allocate_port()` in `_support/ports.py`. That covers owned proxies, scripted upstreams, fake providers, OTLP sinks, database relays and Redis. A server that binds port 0 itself and keeps that socket, such as `wire_server` without a port, needs no allocation. Each pytest process claims a block of 200 ports between 10000 and 32767 by holding an exclusive lock file per block in the system temp directory, so two runs on one machine, such as parallel worktrees, never share a block, and the OS drops the lock when the process exits. The range sits above the fixed service ports and below the Linux and macOS ephemeral ranges, so a port-0 bind or an outbound connection never lands in a block. Ports inside a block go out in order and are never handed out twice in a run, a port that something outside the harness already holds is skipped, and a process that uses up its block claims the next free one. A test that restarts a server on the same port allocates that port up front instead of reading it back from the first server. TQ012 in `scripts/check_test_quality.py` rejects reading a port back from a port-0 bind anywhere under `tests/integration` diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 3c12414553d..fd61ca3aea2 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -150,9 +150,9 @@ def _is_reaped(process: psutil.Process) -> bool: return False -def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seconds: float) -> None: +def _wait_until_reaped(process: psutil.Process, seconds: float) -> None: deadline: Final = time.monotonic() + seconds - while not _is_reaped(process) and not _is_stopped(supervisor): + while not _is_reaped(process): assert time.monotonic() < deadline, f"The supervisor did not reap worker {process.pid} within {seconds:.0f}s" time.sleep(0.05) @@ -160,8 +160,6 @@ def _wait_until_reaped(process: psutil.Process, supervisor: psutil.Process, seco @dataclass(frozen=True, slots=True) class WorkerKill: victim: psutil.Process - supervisor: psutil.Process - reap_seconds: float fired: threading.Event def freeze(self) -> psutil.Process: @@ -173,7 +171,6 @@ class WorkerKill: with suppress(psutil.NoSuchProcess): self.victim.kill() _wait_until_gone(self.victim, STRICT_HEALTHCHECK_SECONDS) - _wait_until_reaped(self.victim, self.supervisor, self.reap_seconds) self.fired.set() @@ -250,22 +247,22 @@ class OwnedProxy: assert pid in before, f"{pid} is not a serving worker of this proxy: {sorted(before)}" return psutil.Process(pid), before - def _worker_kill(self, victim: psutil.Process) -> WorkerKill: - reap_seconds: Final = healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS - return WorkerKill(victim, psutil.Process(self.process.pid), reap_seconds, threading.Event()) + def _replacement_of(self, victim: psutil.Process, before: frozenset[int]) -> int: + _wait_until_reaped(victim, healthcheck_seconds(self.healthcheck) + STRICT_HEALTHCHECK_SECONDS) + return self._await_replacement(before) @contextmanager def replacing_worker(self, pid: int) -> Iterator[WorkerKill]: victim, before = self._victim(pid) - kill: Final = self._worker_kill(victim) + kill: Final = WorkerKill(victim, threading.Event()) yield kill assert kill.fired.is_set(), f"replacing_worker({pid}) exited without killing the worker" - self._await_replacement(before) + self._replacement_of(victim, before) def kill_worker(self, pid: int) -> int: victim, before = self._victim(pid) - self._worker_kill(victim)() - return self._await_replacement(before) + WorkerKill(victim, threading.Event())() + return self._replacement_of(victim, before) @contextmanager def suspended(self, pid: int) -> Iterator[psutil.Process]: diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index b261973973c..e5ce909ab18 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -6,6 +6,7 @@ import time from collections.abc import Callable, Generator, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass +from functools import partial from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType @@ -60,6 +61,7 @@ class Reply: abort_after: int | None = None gate_after_first: threading.Event | None = None pause_between_chunks: float = 0 + heartbeat: bytes | None = None headers: Mapping[str, str] = MappingProxyType({}) drop_connection: bool = False @@ -82,10 +84,14 @@ class Wire: return self.connected.qsize() - self.closed.qsize() -def _await_release(gate: threading.Event, closing: threading.Event) -> bool: +def _await_release( + gate: threading.Event, closing: threading.Event, beat: Callable[[], None] | None = None, every: float = 0.05 +) -> bool: while not closing.is_set(): - if gate.wait(timeout=0.05): + if gate.wait(timeout=every): return True + if beat is not None: + beat() return gate.is_set() @@ -117,6 +123,10 @@ def wire_server( super().finish() closed.put(f"{self.client_address[0]}:{self.client_address[1]}") + def write_chunk(self, chunk: bytes) -> None: + self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + self.wfile.flush() + def respond(self) -> None: headers: Final = {name.lower(): value for name, value in self.headers.items()} try: @@ -156,12 +166,20 @@ def wire_server( for index, chunk in enumerate(reply.chunks): if reply.abort_after == index: break - self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) - self.wfile.flush() + self.write_chunk(chunk) if ( index == 0 and reply.gate_after_first is not None - and not _await_release(reply.gate_after_first, closing) + and not ( + _await_release(reply.gate_after_first, closing) + if reply.heartbeat is None + else _await_release( + reply.gate_after_first, + closing, + partial(self.write_chunk, reply.heartbeat), + reply.pause_between_chunks, + ) + ) ): break if reply.pause_between_chunks and index + 1 < len(reply.chunks): diff --git a/tests/integration/compatibility/test_audio_transcriptions_wire.py b/tests/integration/compatibility/test_audio_transcriptions_wire.py index 7845e3d2423..8f9a83e9912 100644 --- a/tests/integration/compatibility/test_audio_transcriptions_wire.py +++ b/tests/integration/compatibility/test_audio_transcriptions_wire.py @@ -1,15 +1,18 @@ +import io import json import math import mimetypes import uuid +import wave from email.message import Message from email.parser import BytesParser from email.policy import HTTP +from hashlib import sha256 from typing import Final, cast import httpx import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from openai import OpenAI @@ -601,6 +604,31 @@ def test_audio_transcription_consumes_bracketed_form_metadata_tags(gateway: Gate ) == (tag,), rows +@pytest.mark.parametrize( + "identity_param", ("openai_identity_provider_id", "openai_service_account_id", "openai_identity_token_file") +) +def test_audio_transcription_refuses_a_bracketed_identity_form_key_before_the_upstream( + gateway: Gateway, identity_param: str +) -> None: + wav: Final = _WAV_BYTES + uuid.uuid4().bytes + + def respond(request: Request) -> Reply: + return Reply(body=json.dumps(_JSON_TRANSCRIPT).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/whisper-1", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key" + ) + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, "response_format": "json", f"{identity_param}[]": "caller-chosen-" + uuid.uuid4().hex}, + {"file": ("a.wav", wav, _WAV_CONTENT_TYPE)}, + ) + assert response.status_code == 401, response.text + assert identity_param in response.text, response.text + assert wire.drain() == () + + def test_audio_transcription_preserves_the_client_file_content_type(gateway: Gateway) -> None: pytest.skip( "BUG: the proxy discards the client's file part content type and re-derives it from the filename, " @@ -671,3 +699,172 @@ def test_audio_transcription_bills_audio_tokens_at_the_input_rate_when_no_audio_ rel_tol=1e-9, ), response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/audio/transcriptions")] + + +_GPT_TRANSCRIBE_VERBOSE_JSON_REFUSAL: Final = { + "error": { + "message": "response_format 'verbose_json' is not compatible with model 'gpt-transcribe'", + "type": "invalid_request_error", + "param": "response_format", + "code": "unsupported_value", + } +} +_UNDECODABLE_M4A_BYTES: Final = b"\x00\x00\x00\x20ftypM4A \x00\x00\x02\x00M4A isommp42\x00\x00\x00\x08free" +_DEPLOYMENT_COST_PER_SECOND: Final = 0.0003 + + +def _openai_gpt_transcribe_reply(request: Request, reported_seconds: float) -> Reply: + fields: Final = dict(_text_parts(_multipart_parts(request))) + if fields.get("response_format") == "verbose_json": + return Reply(status=400, body=json.dumps(_GPT_TRANSCRIBE_VERBOSE_JSON_REFUSAL).encode()) + return Reply( + body=json.dumps({"text": "hello world", "usage": {"type": "duration", "seconds": reported_seconds}}).encode() + ) + + +def _uploaded_files(request: Request) -> tuple[tuple[str, str, bytes], ...]: + return tuple((name, filename, payload) for name, filename, _, payload in _file_parts(_multipart_parts(request))) + + +def _silent_wav(seconds: float) -> bytes: + output: Final = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(b"\x00\x00" * round(16000 * seconds)) + return output.getvalue() + + +def _advertised_cost_per_second(gateway: Gateway, model: str) -> float: + entries: Final = gateway.get("/v2/model/info", {"model": model})["data"] + assert isinstance(entries, list) and len(entries) == 1, entries + price: Final = object_value(object_value(entries[0])["model_info"])["input_cost_per_second"] + assert isinstance(price, float) and price > 0, entries + return price + + +def _response_cost(response: httpx.Response) -> float: + return float(response.headers["x-litellm-response-cost"]) if "x-litellm-response-cost" in response.headers else 0.0 + + +def _logged_spend(call_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return float(str(rows[0]["spend"])) + + +def test_gpt_transcribe_without_response_format_reaches_openai_without_verbose_json(gateway: Gateway) -> None: + upload: Final = _UNDECODABLE_M4A_BYTES + uuid.uuid4().bytes + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, 2)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-transcribe", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key" + ) + with _sdk(gateway, "/v1", gateway.key) as sdk: + response: Final = sdk.audio.transcriptions.with_raw_response.create( + model=model, file=("a.m4a", upload, "audio/mp4") + ) + requests: Final = wire.drain() + + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["text"] == "hello world", response.text + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _text_parts(_multipart_parts(requests[0])) == (("model", "gpt-transcribe"),) + assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) + + +@pytest.mark.parametrize( + "configured_cost_per_second", + (None, _DEPLOYMENT_COST_PER_SECOND), + ids=("cost_map_price", "deployment_price"), +) +def test_gpt_transcribe_unmeasurable_upload_bills_the_seconds_openai_reports( + gateway: Gateway, configured_cost_per_second: float | None +) -> None: + reported_seconds: Final = 2.0 + upload: Final = _UNDECODABLE_M4A_BYTES + uuid.uuid4().bytes + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, reported_seconds)) as wire, + gateway.scenario() as scenario, + ): + model: Final = ( + scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_second=configured_cost_per_second, + ) + if configured_cost_per_second is not None + else scenario.model( + model="openai/gpt-transcribe", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key" + ) + ) + key: Final = scenario.key(models=[model]) + cost_per_second: Final = ( + configured_cost_per_second + if configured_cost_per_second is not None + else _advertised_cost_per_second(gateway, model) + ) + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, "response_format": "json"}, + {"file": ("a.m4a", upload, "audio/mp4")}, + key=key, + ) + expected_cost: Final = reported_seconds * cost_per_second + + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["text"] == "hello world", response.text + assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) + assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) + key_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', + (sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0, + seconds=70, + ) + assert math.isclose(float(str(key_rows[0]["spend"])), expected_cost), key_rows + requests: Final = wire.drain() + + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) + + +def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_seconds(gateway: Gateway) -> None: + measured_seconds: Final = 1.0 + upload: Final = _silent_wav(measured_seconds) + + with ( + wire_server(lambda request: _openai_gpt_transcribe_reply(request, 7)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_second=_DEPLOYMENT_COST_PER_SECOND, + ) + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model, "response_format": "json"}, + {"file": ("a.wav", upload, _WAV_CONTENT_TYPE)}, + ) + expected_cost: Final = measured_seconds * _DEPLOYMENT_COST_PER_SECOND + + assert response.status_code == 200, response.text + assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) + assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) + requests: Final = wire.drain() + + assert [(request.method, request.target) for request in requests] == [("POST", "/v1/audio/transcriptions")] + assert _uploaded_files(requests[0]) == (("file", "a.wav", upload),) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index c9729353c8e..f98dda5fb0c 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -834,6 +834,61 @@ AZURE_AI_GPT_5_4_MINI_LATEST_STREAM: Final = CostTrackingTestCase( ) +# azure_ai/grok-4.6 (Grok on Azure AI Foundry reports reasoning_tokens outside completion_tokens, like xai) +AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION: Final = CostTrackingTestCase( + name="azure-ai-grok-4.6-reasoning_folded_into_completion", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/grok-4.6", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-grok-4.6-reasoning"}], + "stream": False, + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "model": "grok-4.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2552, + "completion_tokens_details": {"reasoning_tokens": 300}, + }, + }, + ), + expected=ExactExpected( + spend=0.0108, input_cost=0.00368, output_cost=0.00712, prompt_tokens=1840, completion_tokens=712 + ), +) + +AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION_STREAM: Final = CostTrackingTestCase( + name="azure-ai-grok-4.6-reasoning_folded_into_completion-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/grok-4.6", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-grok-4.6-reasoning-stream"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"grok-4.6","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"grok-4.6","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"grok-4.6","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2552,"completion_tokens_details":{"reasoning_tokens":300}}}\n\n', + "data: [DONE]\n\n", + ), + ), + expected=ExactExpected( + spend=0.0108, input_cost=0.00368, output_cost=0.00712, prompt_tokens=1840, completion_tokens=712 + ), +) + + # azure_ai/claude-haiku-4-5 (Azure AI Foundry Anthropic, web search billed per server_tool_use.web_search_requests) AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH: Final = CostTrackingTestCase( name="azure-ai-claude-haiku-4-5-messages_web_search_two_requests", @@ -1978,6 +2033,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_INVOKE_HAIKU_STREAM, AZURE_AI_GPT_5_4_MINI_LATEST, AZURE_AI_GPT_5_4_MINI_LATEST_STREAM, + AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION, + AZURE_AI_GROK_4_6_REASONING_FOLDED_INTO_COMPLETION_STREAM, AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH, AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_STREAM, AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH_ONE_REQUEST, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 5e2b99a113f..2ef42249a7e 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -764,6 +764,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 7.5e-07, "output_cost_per_token": 4.5e-06, }, + "azure_ai/grok-4.6": { + "litellm_provider": "azure_ai", + "mode": "chat", + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + }, "azure_ai/claude-haiku-4-5": { "litellm_provider": "azure_ai", "mode": "chat", diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 2b9ec658d87..9945d16af07 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -1352,3 +1352,40 @@ def test_evaluation_mode_health_check_of_databricks_ai_decide_resolves_the_mode_ }, ) ] + + +def _openai_gpt_transcribe_health_reply(request: Request) -> Reply: + if _multipart_text_fields(_multipart_parts(request)).get("response_format") == "verbose_json": + return Reply( + status=400, + body=json.dumps( + { + "error": { + "message": "response_format 'verbose_json' is not compatible with model 'gpt-transcribe'", + "type": "invalid_request_error", + "param": "response_format", + "code": "unsupported_value", + } + } + ).encode(), + ) + return Reply(body=json.dumps({"text": "healthy", "usage": {"type": "duration", "seconds": 1}}).encode()) + + +def test_health_check_of_gpt_transcribe_probes_without_verbose_json_and_reports_it_healthy(gateway: Gateway) -> None: + with wire_server(_openai_gpt_transcribe_health_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-transcribe", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + model_info={"mode": "audio_transcription"}, + ) + report: Final = _health_report(gateway, model) + requests: Final = wire.drain() + + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert requests, report + assert { + (request.method, request.target, tuple(sorted(_multipart_text_fields(_multipart_parts(request)).items()))) + for request in requests + } == {("POST", "/v1/audio/transcriptions", (("model", "gpt-transcribe"),))} diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 90d204722ac..3e4bb568394 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -53,6 +53,13 @@ SELECT encode(sha256(convert_to(%s || '-' || n::text, 'UTF8')), 'hex'), %s, '{}' FROM generate_series(1, %s::int) AS n """ TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))" +LOCK_MODEL_ROW_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s FOR UPDATE' +LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS: Final = 7 +OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """ +SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age +FROM pg_stat_activity +WHERE %s::int = ANY(pg_blocking_pids(pid)) +""" # Sessions blocked on the advisory lock the given backend holds, and nothing else on the shared rig. WAITERS_ON_HELD_LOCK_SQL: Final = """ SELECT count(*)::int AS waiting @@ -131,6 +138,13 @@ def _waiters_on_lock_held_by(backend_pid: int) -> int: return waiting +def _oldest_transaction_blocked_by(backend_pid: int) -> float: + rows: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(backend_pid),)) + age: Final = rows[0]["age"] + assert isinstance(age, float) + return age + + def _remove_team_by_sql(team_id: str) -> None: """Cleanup for a team the test expects to have deleted itself. Whatever a failed delete left behind (row, memberships, `teams` references) goes by SQL so the shared rig stays clean without sending @@ -161,6 +175,38 @@ def _create_team(scenario: Scenario, **fields: JsonValue) -> str: return team_id +def _create_team_model(scenario: Scenario, team_id: str) -> str: + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": f"integration-{uuid.uuid4().hex}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-integration"}, + "model_info": {"team_id": team_id}, + }, + ) + model_id: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,)) + return model_id + + +def _team_model_rows(model_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,)) + + +def _delete_while_locked(gateway: Gateway, team_id: str, lock_sql: str, lock_target: str) -> httpx.Response: + with ThreadPoolExecutor(max_workers=1) as pool: + with psycopg.connect(os.environ["DATABASE_URL"]) as holder: + holder.execute(lock_sql, (lock_target,)) + pending: Final = _post_in_background(pool, gateway, "/team/delete", {"team_ids": [team_id]}) + eventually( + lambda: _oldest_transaction_blocked_by(holder.info.backend_pid), + lambda age: age > LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS, + seconds=30, + ) + assert not pending.done(), "delete returned while the lock was still held" + return pending.result(timeout=60) + + def _generate_key(scenario: Scenario, **fields: JsonValue) -> str: """A key the team delete is expected to remove; cleanup only deletes it if it is still live.""" key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"]) @@ -392,6 +438,36 @@ def test_delete_waits_for_the_team_advisory_lock_and_completes_after_release( assert len(_tombstones(team)) == 1, _tombstones(team) +def test_delete_completes_after_waiting_on_a_locked_team_model_row_past_the_default_transaction_timeout( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + team: Final = _create_team(scenario) + model_id: Final = _create_team_model(scenario, team) + response: Final = _delete_while_locked(gateway, team, LOCK_MODEL_ROW_SQL, model_id) + assert response.status_code == 200, response.text + assert response.json() == {"deleted_teams": [team]} + assert _team_rows(team) == [] + assert _team_model_rows(model_id) == [] + assert len(_tombstones(team)) == 1, _tombstones(team) + + +def test_delete_completes_after_waiting_on_the_team_lock_past_the_default_transaction_timeout( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + user: Final = scenario.user() + team: Final = _create_team(scenario) + _bulk_member_add(gateway, team, [user]) + response: Final = _delete_while_locked(gateway, team, TAKE_TEAM_LOCK_SQL, team) + assert response.status_code == 200, response.text + assert response.json() == {"deleted_teams": [team]} + assert _team_rows(team) == [] + assert _membership_user_ids(team) == [] + assert _user_teams(user) == [] + assert len(_tombstones(team)) == 1, _tombstones(team) + + def test_deleting_two_teams_sharing_a_member_in_one_call_clears_both_from_the_member(gateway: Gateway) -> None: with gateway.scenario() as scenario: user: Final = scenario.user() diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index cb7f4c2b9ce..e0b2e4a336f 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -40,6 +40,7 @@ FirstAttempt = Literal[ "drop_after_message_start", "error_frame", "error_frame_after_message_start", + "error_frame_after_content", "http_status", "drop_after_content", ] @@ -66,6 +67,8 @@ class _Failure: return status_reply(self.status) case "drop_after_content": return dropping_reply(chunks, abort_after=3) + case "error_frame_after_content": + return stream_reply((*chunks[:3], error_frame(self.status, f"scripted {self.status}"))) _MID_STREAM_FAILURES: Final = ( @@ -124,12 +127,12 @@ def _stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[int, tuple return response.status_code, parse_sse(response.text) +def _rows(request_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) + + def _success_rows(request_id: str) -> list[dict[str, JsonValue]]: - return eventually( - lambda: read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), - lambda rows: len(rows) >= 1, - seconds=70, - ) + return eventually(lambda: _rows(request_id), lambda rows: len(rows) >= 1, seconds=70) def _assert_completed_by_retry( @@ -141,6 +144,8 @@ def _assert_completed_by_retry( assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] * attempts assert _success_rows(_served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] + superseded_rows: Final = {attempt: _rows(_served_id(prompt, attempt)) for attempt in range(1, attempts)} + assert not any(superseded_rows.values()), superseded_rows def _assert_rejected_before_content( @@ -302,6 +307,23 @@ def test_stream_dropping_after_content_is_not_retried(gateway: Gateway) -> None: assert [request.target for request in wire.drain()] == ["/v1/messages"] +def test_stream_error_frame_after_content_is_not_retried_and_keeps_its_spend_row(gateway: Gateway) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("error_frame_after_content", 529), attempts)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(gateway, _body(model, prompt, "deployment")) + assert status == 200, events + assert event_types(events) == (*LIFECYCLE[:3], "error"), events + assert message_id(events) == _served_id(prompt, 1), events + assert delta_text(events) == _TEXT, events + assert [request.target for request in wire.drain()] == ["/v1/messages"] + assert [row["model_group"] for row in _success_rows(_served_id(prompt, 1))] == [model] + + def test_stream_rejected_with_401_before_it_opens_is_not_retried(gateway: Gateway) -> None: prompt: Final = _prompt() attempts: Final = Attempts() diff --git a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py index 2e8a9d0d7e4..15c0593f80d 100644 --- a/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py +++ b/tests/integration/messages_endpoint/test_anthropic_messages_caller_tools_wire.py @@ -364,7 +364,7 @@ def test_spend_log_keeps_the_tools_the_caller_sent(rig: Rig, model: str, guardra assert object_value(rows[0]["proxy_server_request"])["tools"] == [strict_tool()] -@pytest.mark.parametrize("tools", [7, "x" * 5000, ["lint"]], ids=["int", "five_kb", "list_of_str"]) +@pytest.mark.parametrize("tools", [7, "x" * 5000, ["lint"], ""], ids=["int", "five_kb", "list_of_str", "empty_string"]) def test_a_malformed_tools_field_is_answered_with_an_error_and_the_proxy_keeps_serving( rig: Rig, tools: JsonValue ) -> None: @@ -379,14 +379,6 @@ def test_a_malformed_tools_field_is_answered_with_an_error_and_the_proxy_keeps_s assert reply_text(control_status, control_raw, streaming=False) == f"linted {control_nonce.removeprefix('nonce-')}" -def test_an_empty_string_tools_field_reaches_claude_as_sent(rig: Rig) -> None: - request_nonce: Final = nonce() - status, raw = rig.send(messages_body(CLAUDE_MODEL, request_nonce, tools="")) - - assert reply_text(status, raw, streaming=False) == f"linted {request_nonce.removeprefix('nonce-')}" - assert sent_body(rig.received(request_nonce), CLAUDE_TARGET)["tools"] == "" - - @pytest.mark.parametrize( ("model", "target"), [(CLAUDE_MODEL, CLAUDE_TARGET), (OPENAI_MODEL, RESPONSES_TARGET)], diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 9c2a1ecd2c4..587a21e4b62 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -817,7 +817,7 @@ def test_a_failing_bucket_still_bills_the_batch_and_the_other_sink_keeps_its_lin assert [record["custom_id"] for record in _records(gcs_body)] == [custom_id] assert spend > 0 s3_name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id) - assert rig.s3.tried(s3_name) >= 1 + eventually(lambda: rig.s3.tried(s3_name), lambda tries: tries >= 1, POLLER_SECONDS) assert s3_name not in rig.s3.objects() diff --git a/tests/integration/observability/test_generic_guardrail_api_wire.py b/tests/integration/observability/test_generic_guardrail_api_wire.py new file mode 100644 index 00000000000..0796feb5f37 --- /dev/null +++ b/tests/integration/observability/test_generic_guardrail_api_wire.py @@ -0,0 +1,308 @@ +import json +import threading +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway, eventually +from integration._support.openai_wire import chat_reply +from integration._support.ports import allocate_port +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GEMINI_MODEL: Final = "gemini/gemini-3.8-flash" +ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" +OPENAI_MODEL: Final = "openai/gpt-5.6" +ECHO_GEMINI_TOOLS: Final = "echo-gemini-tools" +NONE_CHAT_ROWS: Final = "none-chat-rows" +NONE_CHAT_ROWS_DEFAULT: Final = "none-chat-rows-default" +NONE_ANTHROPIC_IMAGE: Final = "none-anthropic-image" +ECHO_ANTHROPIC_MESSAGES: Final = "echo-anthropic-messages" +FORWARDS_MESSAGES_AS_SENT: Final = frozenset({NONE_CHAT_ROWS, NONE_ANTHROPIC_IMAGE, ECHO_ANTHROPIC_MESSAGES}) +FAIL_OPEN_CONNECT_ERROR: Final = "fail-open-connect-error" +FAIL_OPEN_TIMEOUT: Final = "fail-open-timeout" +FAIL_OPEN_HTTP_STATUS: Final = "fail-open-http-status" +FAIL_OPEN_LINE: Final = "Generic Guardrail API error (fail-open)" +PNG: Final = ( + "iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAIAAACQkWg2AAAAFklEQVR4nGO4IydHEmIY1TCqYfhqAACaMxgQdGu1YAAAAABJRU5ErkJggg==" +) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +def _verdict_none(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _verdict_echo(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + echoed: Final = {key: body[key] for key in ("structured_messages", "tools") if body.get(key)} + return Reply(body=json.dumps({"action": "NONE", **echoed}).encode()) + + +def _unavailable(_request: Request) -> Reply: + return Reply(status=503, body=b'{"error": "guardrail unavailable"}') + + +def _gemini_reply(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1, "totalTokenCount": 4}, + } + ).encode() + ) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_guardrail_wire", "claude-sonnet-5-5", "ok")) + + +def _openai_reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-guardrail-wire", "gpt-5.6", "ok", stream=False) + + +def _guardrail(name: str, api_base: str, **params: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-guardrail-key", + **params, + }, + } + + +@dataclass(frozen=True, slots=True) +class GuardrailRig: + owned: OwnedProxy + policies: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[GuardrailRig]: + directory: Final = tmp_path_factory.mktemp("generic-guardrail-wire") + release: Final = threading.Event() + + def held_until_teardown(request: Request) -> Reply: + release.wait() + return _verdict_none(request) + + with ExitStack() as stack: + policies: Final = MappingProxyType( + { + ECHO_GEMINI_TOOLS: stack.enter_context(wire_server(_verdict_echo)), + NONE_CHAT_ROWS: stack.enter_context(wire_server(_verdict_none)), + NONE_CHAT_ROWS_DEFAULT: stack.enter_context(wire_server(_verdict_none)), + NONE_ANTHROPIC_IMAGE: stack.enter_context(wire_server(_verdict_none)), + ECHO_ANTHROPIC_MESSAGES: stack.enter_context(wire_server(_verdict_echo)), + } + ) + stalled: Final = stack.enter_context(wire_server(held_until_teardown)) + stack.callback(release.set) + unavailable: Final = stack.enter_context(wire_server(_unavailable)) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "generic-guardrail-wire.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + *( + _guardrail(name, wire.url, forward_messages_as_sent=name in FORWARDS_MESSAGES_AS_SENT) + for name, wire in policies.items() + ), + _guardrail(FAIL_OPEN_CONNECT_ERROR, f"http://127.0.0.1:{allocate_port()}", fail_on_error=False), + _guardrail(FAIL_OPEN_TIMEOUT, stalled.url, fail_on_error=False, timeout=1), + _guardrail(FAIL_OPEN_HTTP_STATUS, unavailable.url, fail_on_error=False), + ], + } + ) + ) + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, {}, config=path)) + yield GuardrailRig(owned, policies) + + +def _bodies(wire: Wire) -> list[dict[str, JsonValue]]: + return [_JSON.validate_json(request.body) for request in wire.drain()] + + +def _fail_open_lines(log: Path, guardrail: str) -> list[str]: + return [ + line + for line in log.read_text(errors="replace").splitlines() + if FAIL_OPEN_LINE in line and f"guardrail_name={guardrail} " in line + ] + + +def test_gemini_native_tools_reach_the_guardrail_and_gemini_as_sent(rig: GuardrailRig) -> None: + tools: Final = [ + {"function_declarations": [{"name": "get_weather", "parameters": {"type": "object", "properties": {}}}]}, + {"googleSearch": {}}, + ] + with wire_server(_gemini_reply) as gemini, rig.owned.gateway.scenario() as scenario: + model: Final = scenario.model(model=GEMINI_MODEL, api_base=gemini.url, api_key="synthetic-gemini-key") + response: Final = rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "tools": tools, + "guardrails": [ECHO_GEMINI_TOOLS], + }, + ) + assert response.status_code == 200, response.text + assert [body["tools"] for body in _bodies(rig.policies[ECHO_GEMINI_TOOLS])] == [tools] + sent: Final = _bodies(gemini) + assert len(sent) == 1, sent + assert {"googleSearch": {}} in sent[0]["tools"], sent + assert [tool for tool in sent[0]["tools"] if "type" in tool] == [], sent + + +CHAT_ROWS: Final = ( + {"role": "user", "name": "pat", "content": "weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"index": 0, "id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, +) +OPENAI_CHAT_ROWS: Final = ( + {"role": "user", "content": "weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, +) + + +@pytest.mark.parametrize( + ("guardrail", "seen"), + [(NONE_CHAT_ROWS, CHAT_ROWS), (NONE_CHAT_ROWS_DEFAULT, OPENAI_CHAT_ROWS)], + ids=["forward_messages_as_sent", "default"], +) +def test_guardrail_sees_chat_rows_as_sent_only_when_configured( + rig: GuardrailRig, guardrail: str, seen: tuple[dict[str, JsonValue], ...] +) -> None: + messages: Final = list(CHAT_ROWS) + with wire_server(_openai_reply) as openai_peer, rig.owned.gateway.scenario() as scenario: + model: Final = scenario.model( + model=OPENAI_MODEL, api_base=f"{openai_peer.url}/v1", api_key="synthetic-openai-key" + ) + response: Final = rig.owned.gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": messages, "guardrails": [guardrail]} + ) + assert response.status_code == 200, response.text + assert [body["structured_messages"] for body in _bodies(rig.policies[guardrail])] == [list(seen)] + assert len(openai_peer.drain()) == 1 + + +def test_guardrail_sees_an_anthropic_image_block_on_chat_completions_as_sent(rig: GuardrailRig) -> None: + image: Final = [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG}}] + with wire_server(_anthropic_reply) as anthropic_peer, rig.owned.gateway.scenario() as scenario: + model: Final = scenario.model( + model=ANTHROPIC_MODEL, api_base=anthropic_peer.url, api_key="synthetic-anthropic-key" + ) + rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": image}, {"role": "user", "content": "what color?"}], + "guardrails": [NONE_ANTHROPIC_IMAGE], + }, + ) + seen: Final = _bodies(rig.policies[NONE_ANTHROPIC_IMAGE]) + assert [body["structured_messages"][0]["content"] for body in seen] == [image], seen + + +def test_echoing_guardrail_leaves_the_anthropic_messages_request_unchanged(rig: GuardrailRig) -> None: + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What color is this square? Also check the weather in Paris."}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG}}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "Trip notes."}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny, 21C"}, + {"type": "text", "text": "Answer both in one sentence."}, + ], + }, + ] + tools: Final = [ + { + "name": "get_weather", + "description": "Weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + with wire_server(_anthropic_reply) as anthropic_peer, rig.owned.gateway.scenario() as scenario: + model: Final = scenario.model( + model=ANTHROPIC_MODEL, api_base=anthropic_peer.url, api_key="synthetic-anthropic-key" + ) + response: Final = rig.owned.gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": messages, + "tools": tools, + "guardrails": [ECHO_ANTHROPIC_MESSAGES], + }, + ) + assert response.status_code == 200, response.text + assert len(rig.policies[ECHO_ANTHROPIC_MESSAGES].drain()) == 1 + assert [body["messages"] for body in _bodies(anthropic_peer)] == [messages] + + +@pytest.mark.parametrize( + ("guardrail", "tokens"), + [ + (FAIL_OPEN_CONNECT_ERROR, ("error_type=ConnectError",)), + (FAIL_OPEN_TIMEOUT, ("error_type=Timeout",)), + (FAIL_OPEN_HTTP_STATUS, ("error_type=HTTPStatusError", "http_status_code=503")), + ], + ids=["connect_error", "timeout", "http_status"], +) +def test_fail_open_log_line_names_the_error_type(rig: GuardrailRig, guardrail: str, tokens: tuple[str, ...]) -> None: + with wire_server(_openai_reply) as openai_peer, rig.owned.gateway.scenario() as scenario: + model: Final = scenario.model( + model=OPENAI_MODEL, api_base=f"{openai_peer.url}/v1", api_key="synthetic-openai-key" + ) + response: Final = rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "hi"}], "guardrails": [guardrail]}, + ) + assert response.status_code == 200, response.text + assert len(openai_peer.drain()) == 1 + lines: Final = eventually(lambda: _fail_open_lines(rig.owned.log, guardrail), bool, seconds=30) + assert [line for line in lines if not all(token in line for token in tokens)] == [], lines diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py index 01c1f563be4..d845d332237 100644 --- a/tests/integration/observability/test_s3_v2_upload_fanout.py +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -420,7 +420,16 @@ def test_s3_v2_object_rejected_with_a_bucket_wide_code_is_delivered_once_the_fau def test_s3_v2_terminal_object_is_put_once_and_dropped_by_default(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "s3toolarge" + uuid.uuid4().hex[:8] sink: Final = RejectingS3Sink(reject_marker=f"{marker}-huge", reject_status=400, reject_code="EntityTooLarge") - with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + flushing: Final = threading.Event() + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + if f"{marker}-held".encode() in request.body: + flushing.set() + assert release.wait(timeout=60), "the held flush was never released" + return sink.respond(request) + + with wire_server(_chat_reply) as provider, wire_server(respond) as bucket: config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": False}) with ( owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config) as candidate, @@ -428,11 +437,16 @@ def test_s3_v2_terminal_object_is_put_once_and_dropped_by_default(gateway: Gatew ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") key: Final = scenario.key(models=[model]) + _send(candidate, model, key, f"{marker}-held") + eventually(flushing.is_set, bool, seconds=30) _send(candidate, model, key, f"{marker}-huge") - _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-sibling") + _send(candidate, model, key, f"{marker}-sibling") + release.set() + eventually(sink.landed_ids, lambda landed: f"{marker}-sibling" in landed, seconds=60) _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-second-flush") _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-third-flush") - assert sum(1 for r in provider.drain() if r.method == "POST") == 4 + assert sum(1 for r in provider.drain() if r.method == "POST") == 5 + assert sink.landed_ids() == {f"{marker}-{name}" for name in ("held", "sibling", "second-flush", "third-flush")} assert sink.rejected_attempts == 1, ( f"an EntityTooLarge object was PUT {sink.rejected_attempts} times next to delivered siblings; " "with the default s3_drop_on_terminal_error it must be attempted once and dropped" diff --git a/tests/integration/providers/test_azure_reasoning_effort_none_wire.py b/tests/integration/providers/test_azure_reasoning_effort_none_wire.py new file mode 100644 index 00000000000..c95e497c672 --- /dev/null +++ b/tests/integration/providers/test_azure_reasoning_effort_none_wire.py @@ -0,0 +1,58 @@ +from typing import Final + +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_DEPLOYMENT: Final = "gpt-6.1-sol" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "no reasoning please"}] + + +def _reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-effort-none", _DEPLOYMENT, "answered without reasoning", stream=False) + + +def test_reasoning_effort_none_is_rejected_by_litellm_before_reaching_azure(gateway: Gateway) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version="2025-04-01-preview", + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": _MESSAGES, "reasoning_effort": "none"}, + key=scenario.key(), + ) + + assert response.status_code == 400, response.text + assert "reasoning_effort='none'" in response.text, response.text + assert wire.drain() == () + + +def test_reasoning_effort_none_is_dropped_when_the_deployment_drops_params(gateway: Gateway) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version="2025-04-01-preview", + drop_params=True, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": _MESSAGES, "reasoning_effort": "none"}, + key=scenario.key(), + ) + + assert response.status_code == 200, response.text + assert "answered without reasoning" in response.text, response.text + received: Final = wire.drain() + assert [request.method for request in received] == ["POST"], [request.target for request in received] + body: Final = _JSON_OBJECT.validate_json(received[0].body) + assert "reasoning_effort" not in body and "reasoning" not in body, body diff --git a/tests/integration/providers/test_azure_responses_api_version_wire.py b/tests/integration/providers/test_azure_responses_api_version_wire.py new file mode 100644 index 00000000000..13c2f48dee2 --- /dev/null +++ b/tests/integration/providers/test_azure_responses_api_version_wire.py @@ -0,0 +1,108 @@ +from typing import Final +from urllib.parse import parse_qs, urlsplit + +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_DEPLOYMENT: Final = "gpt-5.4-nano" +_DATED_GA_API_VERSION: Final = "2024-10-21" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _reply(request: Request) -> Reply: + if urlsplit(request.target).path.endswith("/responses"): + return responses_reply("resp_azure_version", _DEPLOYMENT, "azure answered", stream=False) + return chat_reply("chatcmpl-azure-version", _DEPLOYMENT, "azure answered", stream=False) + + +def _dated_azure_model(scenario: Scenario, wire: Wire, model: str, responses_api_version: str | None = None) -> str: + if responses_api_version is None: + return scenario.model( + model=model, api_base=wire.url, api_key="synthetic-azure-key", api_version=_DATED_GA_API_VERSION + ) + return scenario.model( + model=model, + api_base=wire.url, + api_key="synthetic-azure-key", + api_version=_DATED_GA_API_VERSION, + responses_api_version=responses_api_version, + ) + + +def _only_post(wire: Wire) -> tuple[str, dict[str, str], dict[str, JsonValue]]: + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + target: Final = urlsplit(posts[0].target) + query: Final = {name: values[-1] for name, values in parse_qs(target.query).items()} + return target.path, query, _JSON_OBJECT.validate_json(posts[0].body) + + +def test_chat_keeps_the_dated_api_version_and_never_sends_responses_api_version(gateway: Gateway) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = _dated_azure_model(scenario, wire, f"azure/{_DEPLOYMENT}", responses_api_version="preview") + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "dated chat"}]}, + key=scenario.key(), + ) + + assert response.status_code == 200, response.text + path, query, body = _only_post(wire) + assert (path, query) == ( + f"/openai/deployments/{_DEPLOYMENT}/chat/completions", + {"api-version": _DATED_GA_API_VERSION}, + ) + assert body == { + "model": _DEPLOYMENT, + "messages": [{"role": "user", "content": "dated chat"}], + "stream": False, + }, body + + +def test_responses_on_a_dated_ga_deployment_reaches_the_v1_responses_route(gateway: Gateway) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = _dated_azure_model(scenario, wire, f"azure/{_DEPLOYMENT}") + response: Final = gateway.request( + "POST", "/v1/responses", {"model": model, "input": "dated responses"}, key=scenario.key() + ) + + assert response.status_code == 200, response.text + path, query, body = _only_post(wire) + assert (path, query) == ("/openai/v1/responses", {"api-version": "preview"}) + assert body["input"] == "dated responses", body + + +def test_responses_api_version_overrides_the_chat_api_version_and_stays_out_of_the_body(gateway: Gateway) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = _dated_azure_model( + scenario, wire, f"azure/{_DEPLOYMENT}", responses_api_version="2025-04-01-preview" + ) + response: Final = gateway.request( + "POST", "/v1/responses", {"model": model, "input": "override responses"}, key=scenario.key() + ) + + assert response.status_code == 200, response.text + path, query, body = _only_post(wire) + assert (path, query) == ("/openai/responses", {"api-version": "2025-04-01-preview"}) + assert "responses_api_version" not in body and "api_version" not in body, body + + +def test_chat_bridged_through_azure_responses_on_a_dated_ga_deployment_reaches_v1_responses( + gateway: Gateway, +) -> None: + with wire_server(_reply) as wire, gateway.scenario() as scenario: + model: Final = _dated_azure_model(scenario, wire, f"azure/responses/{_DEPLOYMENT}") + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "bridged chat"}]}, + key=scenario.key(), + ) + + assert response.status_code == 200, response.text + assert "azure answered" in response.text, response.text + path, query, _ = _only_post(wire) + assert (path, query) == ("/openai/v1/responses", {"api-version": "preview"}) diff --git a/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py b/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py new file mode 100644 index 00000000000..ca22ab4f23d --- /dev/null +++ b/tests/integration/providers/test_bedrock_invoke_dangerous_tool_use_beta_wire.py @@ -0,0 +1,138 @@ +import json +from collections.abc import Callable, Mapping +from typing import Final +from urllib.parse import unquote + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +CLAUDE_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +MANTLE_CLAUDE_ID: Final = "anthropic.claude-haiku-4-5" +TOKEN: Final = "synthetic-bedrock-bearer" +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" +KEPT_BETA: Final = "context-1m-2025-08-07" +SAFEGUARDS: Final[JsonValue] = [{"type": "dangerous_tool_use"}] +NO_CACHE: Final[JsonValue] = {"no-cache": True} +INVOKE_TARGET: Final = f"/model/{CLAUDE_ID}/invoke" +MANTLE_TARGET: Final = "/anthropic/v1/messages" +REPLY: Final = json.dumps( + { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": CLAUDE_ID, + "content": [{"type": "text", "text": "scripted reply"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } +).encode() +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_BETAS: Final = TypeAdapter(tuple[str, ...]) + + +def _peer(request: Request) -> Reply: + if unquote(request.target) not in (INVOKE_TARGET, MANTLE_TARGET): + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + assert request.headers["authorization"] == f"Bearer {TOKEN}", request.headers + return Reply(body=REPLY) + + +def _invoke_body_betas(request: Request) -> frozenset[str]: + return frozenset(_BETAS.validate_python(_OBJECT.validate_json(request.body).get("anthropic_beta") or ())) + + +def _mantle_header_betas(request: Request) -> frozenset[str]: + return frozenset(beta.strip() for beta in request.headers.get("anthropic-beta", "").split(",") if beta.strip()) + + +def _delivered(gateway: Gateway, wire: Wire, path: str, body: Mapping[str, JsonValue], client_betas: str) -> Request: + response: Final = gateway.request("POST", path, body, headers={"anthropic-beta": client_betas}) + assert response.status_code == 200, response.text + (sent,) = wire.drain() + return sent + + +def _messages_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "cache": NO_CACHE, + **extra, + } + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": "hi"}], "cache": NO_CACHE, **extra} + + +def _responses_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "input": "hi", "cache": NO_CACHE, **extra} + + +def test_invoke_messages_without_safeguards_sends_the_betas_of_a_request_that_never_had_the_beta( + gateway: Gateway, +) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + with_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + without_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), KEPT_BETA) + ) + assert KEPT_BETA in without_beta, without_beta + assert with_beta == without_beta, (with_beta, without_beta) + + +def test_invoke_messages_with_safeguards_keeps_the_dangerous_beta(gateway: Gateway) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + sent: Final = _invoke_body_betas( + _delivered( + gateway, + wire, + "/v1/messages", + _messages_body(model, safeguards=SAFEGUARDS), + f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}", + ) + ) + assert {KEPT_BETA, DANGEROUS_TOOL_USE_BETA} <= sent, sent + + +@pytest.mark.parametrize( + ("path", "build"), + [("/v1/chat/completions", _chat_body), ("/v1/responses", _responses_body)], + ids=["chat_completions", "responses"], +) +def test_invoke_chat_shaped_routes_without_safeguards_drop_only_the_dangerous_beta( + gateway: Gateway, path: str, build: Callable[[str], dict[str, JsonValue]] +) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + with_beta: Final = _invoke_body_betas( + _delivered(gateway, wire, path, build(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + without_beta: Final = _invoke_body_betas(_delivered(gateway, wire, path, build(model), KEPT_BETA)) + assert KEPT_BETA in without_beta, without_beta + assert with_beta == without_beta, (with_beta, without_beta) + + +def test_mantle_messages_without_safeguards_keeps_the_dangerous_beta(gateway: Gateway) -> None: + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock_mantle/{MANTLE_CLAUDE_ID}", api_key=TOKEN, aws_region_name="us-east-1", api_base=wire.url + ) + sent: Final = _mantle_header_betas( + _delivered(gateway, wire, "/v1/messages", _messages_body(model), f"{KEPT_BETA},{DANGEROUS_TOOL_USE_BETA}") + ) + assert DANGEROUS_TOOL_USE_BETA in sent, sent diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 5179514b2b8..701d185b241 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1147,8 +1147,8 @@ def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( @pytest.mark.parametrize( "tools", - ["", "x" * 5120, ["get_weather"]], - ids=["empty_string", "5kb_string", "list_of_strings"], + [5, "", "x" * 5120, ["get_weather"]], + ids=["int", "empty_string", "5kb_string", "list_of_strings"], ) def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_peer( counting_proxy: OwnedProxy, mantle_port: int, tools: JsonValue @@ -1161,8 +1161,8 @@ def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_pe ): model: Final = _deployment(scenario, runtime.url) refused: Final = _count(gateway, {**_bare(model), "tools": tools}) - assert 400 <= refused.status_code < 600, refused.text - assert "input_tokens" not in refused.text, refused.text + assert refused.status_code == 400, refused.text + assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["type"] == "invalid_request_error" assert runtime.drain() == () and mantle.drain() == () assert _generated_then_counted(gateway.client, gateway.key, model, _bare(model)) == (200, 200, _MANTLE_COUNT) assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -1185,8 +1185,13 @@ def test_messages_count_tokens_without_messages_is_rejected( assert runtime.drain() == () and mantle.drain() == () +@pytest.mark.parametrize( + "fields", + [{}, {"tools": 5}, {"web_search_options": None}], + ids=["well_formed", "int_tools", "null_web_search_options"], +) def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer( - counting_proxy: OwnedProxy, mantle_port: int + counting_proxy: OwnedProxy, mantle_port: int, fields: dict[str, JsonValue] ) -> None: gateway: Final = counting_proxy.gateway with ( @@ -1196,7 +1201,7 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request( - "POST", "/v1/messages/count_tokens", _full(model), key="sk-not-a-key-this-proxy-issued" + "POST", "/v1/messages/count_tokens", {**_full(model), **fields}, key="sk-not-a-key-this-proxy-issued" ) assert response.status_code == 401, response.text assert "input_tokens" not in response.text, response.text diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py index 0860d79459e..54923fe0c97 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py @@ -27,6 +27,7 @@ BAD_KEY: Final = "sk-synthetic-bad-key" NATIVE_TARGET: Final = "/openai/v1/chat/completions" CONVERSE_TARGET: Final = f"/model/{GPT}/converse" LONG_VERSION_GPT: Final = "openai.gpt-" + "1" * 30000 +LONG_UNVERSIONED_GPT: Final = "openai.gpt-x" + "1" * 30000 PNG_DATA_URL: Final = ( "data:image/png;base64," "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==" @@ -258,11 +259,12 @@ def test_thirty_thousand_digit_version_is_classified_quickly_and_served_by_conve marker: Final = uuid.uuid4().hex control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool: - control_model: Final = _deployment(scenario, wire) + control_model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_UNVERSIONED_GPT}") model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") control: Final = timed(lambda: _chat(gateway, control_model, control_marker)) assert _content(control.value) == answer(control_marker), control.value.text - _only_request(wire, control_marker) + control_request: Final = _only_request(wire, control_marker) + assert target_of(control_request) == f"/model/{LONG_UNVERSIONED_GPT}/converse", control_request.target liveliness: Final = pool.submit(_timed_liveliness, gateway) served: Final = timed(lambda: _chat(gateway, model, marker)) response: Final = served.value diff --git a/tests/integration/providers/test_malformed_tools_chaos.py b/tests/integration/providers/test_malformed_tools_chaos.py new file mode 100644 index 00000000000..0ab50a5afbe --- /dev/null +++ b/tests/integration/providers/test_malformed_tools_chaos.py @@ -0,0 +1,619 @@ +from __future__ import annotations + +import json +import os +import re +import threading +import uuid +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal, TypeAlias + +import httpx +import psutil +import pytest +from integration._support.anthropic_sse import error_body, message_json, message_stream +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import write_rows +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +_OPENAI_BACKEND: Final = "gpt-4o-mini" +_CLAUDE_BACKEND: Final = "claude-haiku-4-5" +_GEMINI_BACKEND: Final = "gemini-2.5-flash" +_PROVIDER_KEY: Final = "scripted-malformed-tools-chaos-key" +_OWNED_CHAT: Final = "malformed-tools-owned-chat" +_OWNED_STORE: Final = "vs_malformed_tools_owned_store" +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60 +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_NOT_A_LIST: Final = "tools: Input should be a valid list" +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_CHAT_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, +} +_CLAUDE_TOOL: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "description": "Look up the weather", + "input_schema": _WEATHER_SCHEMA, +} +_RESPONSES_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, +} +_GEMINI_TOOL: Final[dict[str, JsonValue]] = { + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA} + ] +} +Kind: TypeAlias = Literal["valid", "rejected", "dropped"] + + +def _marker(label: str) -> str: + return f"{label} marker-{uuid.uuid4().hex}" + + +def _tag(text: str) -> str: + found: Final = _MARKER.search(text) + return found.group(1) if found else "none" + + +def _answer(tag: str) -> str: + return f"answer {tag}" + + +def _gemini_body(tag: str) -> bytes: + return json.dumps( + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": _answer(tag)}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11}, + "modelVersion": _GEMINI_BACKEND, + } + ).encode() + + +def _scripted(request: Request) -> Reply: + text: Final = request.body.decode() + tag: Final = _tag(text) + path: Final = request.target.split("?")[0] + stream: Final = bool(text) and _JSON_OBJECT.validate_json(request.body).get("stream") is True + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/responses": + return responses_reply(f"resp_{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/messages" and _JSON_OBJECT.validate_json(request.body).get("tools") == 5: + return Reply(status=400, body=error_body(400, _NOT_A_LIST)) + if path == "/v1/messages" and stream: + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag)) + ) + if path == "/v1/messages": + return Reply(body=message_json(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag))) + if path.endswith(":streamGenerateContent"): + return Reply(content_type="text/event-stream", chunks=(b"data: " + _gemini_body(tag) + b"\r\n\r\n",)) + if path.endswith(":generateContent"): + return Reply(body=_gemini_body(tag)) + return Reply(status=404, body=json.dumps({"error": f"unscripted {request.target}"}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Switchboard: + down: threading.Event + holding: threading.Event + release: threading.Event + held: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + if self.down.is_set(): + return Reply(drop_connection=True) + if self.holding.is_set(): + self.held.put(request.target) + assert self.release.wait(timeout=120), "Held request was never released" + return _scripted(request) + + def reset(self) -> None: + self.down.clear() + self.release.set() + self.holding.clear() + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + wire: Wire + switchboard: _Switchboard + chat: str + claude: str + gemini: str + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + switchboard: Final = _Switchboard(threading.Event(), threading.Event(), threading.Event(), SimpleQueue()) + with ( + gateway_from_environment() as gateway, + wire_server(switchboard.respond) as wire, + gateway.scenario() as scenario, + ): + built: Final = _Rig( + gateway=gateway, + wire=wire, + switchboard=switchboard, + chat=scenario.model(api_base=f"{wire.url}/v1", api_key=_PROVIDER_KEY), + claude=scenario.model(model=f"anthropic/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + gemini=scenario.model(model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + ) + for group in (built.chat, built.claude, built.gemini): + wait_for_worker_readiness(gateway, group, 1, _PROXY_WORKERS) + try: + yield built + finally: + switchboard.reset() + + +@pytest.fixture(autouse=True) +def _reset(rig: _Rig) -> Iterator[None]: + rig.switchboard.reset() + rig.wire.drain() + yield + rig.switchboard.reset() + + +@dataclass(frozen=True, slots=True) +class _Call: + path: str + body: dict[str, JsonValue] + marker: str + kind: Kind + + +def _tools_for(kind: Kind, tool: dict[str, JsonValue]) -> dict[str, JsonValue]: + if kind == "rejected": + return {"tools": 5} + if kind == "dropped": + return {"tools": [tool, 5], "drop_params": True} + return {"tools": [tool]} + + +def _call(rig: _Rig, index: int, kind: Kind, label: str) -> _Call: + marker: Final = _marker(label) + stream: Final = (index // 4) % 2 == 1 + family: Final = index % 4 + if family == 0: + return _Call( + "/v1/chat/completions", + { + "model": rig.chat, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + **_tools_for(kind, _CHAT_TOOL), + }, + marker, + kind, + ) + if family == 1: + return _Call( + "/v1/messages", + { + "model": rig.claude, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + **_tools_for(kind, _CLAUDE_TOOL), + }, + marker, + kind, + ) + if family == 2: + return _Call( + "/v1/responses", + {"model": rig.chat, "input": marker, "stream": stream, **_tools_for(kind, _RESPONSES_TOOL)}, + marker, + kind, + ) + action: Final = "streamGenerateContent?alt=sse" if stream else "generateContent" + return _Call( + f"/v1beta/models/{rig.gemini}:{action}", + {"contents": [{"role": "user", "parts": [{"text": marker}]}], **_tools_for(kind, _GEMINI_TOOL)}, + marker, + kind, + ) + + +def _kind(index: int) -> Kind: + kinds: Final[tuple[Kind, ...]] = ("valid", "rejected", "dropped") + return kinds[index % 3] + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=120, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, call: _Call) -> httpx.Response: + return client.post(call.path, json=call.body, headers={"Authorization": f"Bearer {key}"}) + + +def _status_or_dropped(client: httpx.Client, key: str, call: _Call) -> int | None: + try: + return _post(client, key, call).status_code + except httpx.TransportError: + return None + + +def _deliveries(wire: Wire) -> tuple[str, ...]: + return tuple(request.body.decode() for request in wire.drain()) + + +def _seen(deliveries: tuple[str, ...], marker: str) -> int: + return sum(marker in body for body in deliveries) + + +def _delivered_tool_names(deliveries: tuple[str, ...], marker: str) -> tuple[str, ...]: + (body,) = tuple(body for body in deliveries if marker in body) + return tuple(re.findall(r'"name": ?"([a-z_]+)"', body)) + + +def _assert_wave(calls: tuple[_Call, ...], statuses: tuple[int | None, ...], deliveries: tuple[str, ...]) -> None: + expected_statuses: Final = tuple(400 if call.kind == "rejected" else 200 for call in calls) + assert statuses == expected_statuses, tuple(zip(statuses, (call.marker for call in calls), strict=True)) + seen: Final = tuple(_seen(deliveries, call.marker) for call in calls) + assert seen == tuple(0 if call.kind == "rejected" else 1 for call in calls), seen + names: Final = tuple(_delivered_tool_names(deliveries, call.marker) for call in calls if call.kind != "rejected") + assert all("get_weather" in found for found in names), names + + +def test_mixed_burst_rejects_malformed_tools_and_forwards_the_rest_once(rig: _Rig) -> None: + with ExitStack() as stack: + calls: Final = tuple(_call(rig, index, _kind(index), "burst") for index in range(40)) + clients: Final = _clients(stack, str(rig.gateway.client.base_url), len(calls)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(calls))) + statuses: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), zip(clients, calls, strict=True) + ) + ) + _assert_wave(calls, statuses, _deliveries(rig.wire)) + + +def test_upstream_outage_still_rejects_malformed_tools_and_recovers(rig: _Rig) -> None: + with ExitStack() as stack: + clients: Final = _clients(stack, str(rig.gateway.client.base_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + + def wave(label: str) -> tuple[tuple[_Call, ...], tuple[int | None, ...]]: + calls: Final = tuple(_call(rig, index, _kind(index), label) for index in range(len(clients))) + statuses: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), zip(clients, calls, strict=True) + ) + ) + return calls, statuses + + before_calls, before_statuses = wave("before outage") + _assert_wave(before_calls, before_statuses, _deliveries(rig.wire)) + rig.switchboard.down.set() + outage_calls, outage_statuses = wave("during outage") + for call, status in zip(outage_calls, outage_statuses, strict=True): + assert (status == 400) if call.kind == "rejected" else (status is not None and status >= 500), ( + call.kind, + status, + ) + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + rig.wire.drain() + rig.switchboard.down.clear() + after_calls, after_statuses = wave("after outage") + _assert_wave(after_calls, after_statuses, _deliveries(rig.wire)) + + +def test_held_upstream_does_not_delay_malformed_tools_rejections(rig: _Rig) -> None: + with ExitStack() as stack: + held_calls: Final = tuple(_call(rig, index * 4, "valid", "held") for index in range(8)) + rejected_calls: Final = tuple(_call(rig, index, "rejected", "rejected") for index in range(8)) + clients: Final = _clients(stack, str(rig.gateway.client.base_url), len(held_calls) + len(rejected_calls)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + rig.switchboard.release.clear() + rig.switchboard.holding.set() + stack.callback(rig.switchboard.release.set) + pending: Final[tuple[Future[int | None], ...]] = tuple( + pool.submit(_status_or_dropped, client, rig.gateway.key, call) + for client, call in zip(clients[: len(held_calls)], held_calls, strict=True) + ) + eventually(rig.switchboard.held.qsize, lambda size: size == len(held_calls), seconds=60) + rejected: Final = tuple( + pool.map( + lambda pair: _status_or_dropped(pair[0], rig.gateway.key, pair[1]), + zip(clients[len(held_calls) :], rejected_calls, strict=True), + ) + ) + assert rejected == (400,) * len(rejected_calls), rejected + assert not any(future.done() for future in pending), "a held request finished before its release" + rig.switchboard.release.set() + assert tuple(future.result(timeout=120) for future in pending) == (200,) * len(held_calls) + deliveries: Final = _deliveries(rig.wire) + assert tuple(_seen(deliveries, call.marker) for call in held_calls) == (1,) * len(held_calls) + assert tuple(_seen(deliveries, call.marker) for call in rejected_calls) == (0,) * len(rejected_calls) + + +@dataclass(frozen=True, slots=True) +class _Owned: + proxy: OwnedProxy + wire: Wire + held: SimpleQueue[str] + holding: threading.Event + release: threading.Event + blocked_tool: str + denied_tool: str + + +def _owned_config(path: Path, upstream: str, denied_tool: str) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": _OWNED_CHAT, + "litellm_params": { + "model": f"openai/{_OPENAI_BACKEND}", + "api_base": f"{upstream}/v1", + "api_key": _PROVIDER_KEY, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "litellm_settings": {"drop_params": True}, + "router_settings": {"disable_cooldowns": True}, + "guardrails": [ + { + "guardrail_name": "malformed-tools-policy", + "litellm_params": {"guardrail": "tool_policy", "mode": "pre_call", "default_on": True}, + }, + { + "guardrail_name": "malformed-tools-permission", + "litellm_params": { + "guardrail": "tool_permission", + "mode": "pre_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": "block", + "rules": [{"id": "deny_tool", "tool_name": f"^{denied_tool}$", "decision": "deny"}], + }, + }, + ], + "vector_store_registry": [ + { + "vector_store_name": "malformed-tools-owned-store", + "litellm_params": { + "vector_store_id": _OWNED_STORE, + "custom_llm_provider": "openai", + "api_base": upstream, + "api_key": _PROVIDER_KEY, + }, + } + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def owned(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + suffix: Final = uuid.uuid4().hex[:12] + blocked_tool: Final = f"wipe_disk_{suffix}" + denied_tool: Final = f"run_shell_{suffix}" + held: Final[SimpleQueue[str]] = SimpleQueue() + holding: Final = threading.Event() + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + if holding.is_set(): + held.put(request.target) + assert release.wait(timeout=120), "Held request was never released" + return _scripted(request) + + directory: Final = tmp_path_factory.mktemp("malformed-tools-owned") + with ExitStack() as stack: + gateway: Final = stack.enter_context(gateway_from_environment()) + policy: Final = gateway.request( + "POST", "/v1/tool/policy", {"tool_name": blocked_tool, "input_policy": "blocked"} + ) + assert policy.status_code == 200, policy.text + stack.callback(write_rows, 'DELETE FROM "LiteLLM_ToolTable" WHERE tool_name = %s', (blocked_tool,)) + wire: Final = stack.enter_context(wire_server(respond)) + stack.callback(release.set) + proxy: Final = stack.enter_context( + owned_proxy_process( + gateway, + directory, + {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _PROVIDER_KEY}, + config=_owned_config(directory / "malformed-tools.yaml", wire.url, denied_tool), + workers=2, + ) + ) + yield _Owned(proxy, wire, held, holding, release, blocked_tool, denied_tool) + + +def _owned_chat(owned: _Owned, marker: str, **fields: JsonValue) -> httpx.Response: + return owned.proxy.gateway.request( + "POST", + "/v1/chat/completions", + {"model": _OWNED_CHAT, "messages": [{"role": "user", "content": marker}], **fields}, + ) + + +def _named_tool(name: str) -> dict[str, JsonValue]: + return {"type": "function", "function": {"name": name, "parameters": _WEATHER_SCHEMA}} + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_global_drop_params_wins_over_an_explicit_request_false(owned: _Owned) -> None: + marker: Final = _marker("global drop") + response: Final = _owned_chat(owned, marker, tools=5, drop_params=False) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert "tools" not in _JSON_OBJECT.validate_json(sent.body), sent.body + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_tool_policy_still_blocks_a_tool_kept_beside_a_dropped_entry(owned: _Owned) -> None: + marker: Final = _marker("tool policy") + response: Final = _owned_chat(owned, marker, tools=[_named_tool(owned.blocked_tool), 5]) + assert response.status_code == 400, response.text + assert "Violated tool policy" in response.text and owned.blocked_tool in response.text, response.text + assert tuple(request for request in owned.wire.drain() if marker in request.body.decode()) == () + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_tool_permission_still_denies_a_tool_kept_beside_a_dropped_entry(owned: _Owned) -> None: + marker: Final = _marker("tool permission") + response: Final = _owned_chat( + owned, marker, tools=[_named_tool(owned.denied_tool), "x"], guardrails=["malformed-tools-permission"] + ) + assert response.status_code == 400, response.text + assert "Violated guardrail policy" in response.text and owned.denied_tool in response.text, response.text + assert tuple(request for request in owned.wire.drain() if marker in request.body.decode()) == () + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_anthropic_pass_through_forwards_malformed_tools_unchanged(owned: _Owned) -> None: + marker: Final = _marker("pass-through") + body: Final[dict[str, JsonValue]] = { + "model": _CLAUDE_BACKEND, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "tools": 5, + } + response: Final = owned.proxy.gateway.request("POST", "/anthropic/v1/messages", body) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_json(response.content) == _JSON_OBJECT.validate_json(error_body(400, _NOT_A_LIST)) + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert sent.target == "/v1/messages", sent.target + assert _JSON_OBJECT.validate_json(sent.body) == body, sent.body + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_anthropic_pass_through_streams_well_formed_tools(owned: _Owned) -> None: + marker: Final = _marker("pass-through stream") + body: Final[dict[str, JsonValue]] = { + "model": _CLAUDE_BACKEND, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "tools": [_CLAUDE_TOOL], + "stream": True, + } + response: Final = owned.proxy.gateway.request("POST", "/anthropic/v1/messages", body) + assert response.status_code == 200, response.text + assert response.content == b"".join(message_stream(f"msg_{_tag(marker)}", _CLAUDE_BACKEND, _answer(_tag(marker)))) + (sent,) = tuple(request for request in owned.wire.drain() if marker in request.body.decode()) + assert _JSON_OBJECT.validate_json(sent.body) == body, sent.body + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +def _owned_call(index: int, label: str) -> _Call: + marker: Final = _marker(label) + tools: Final[JsonValue] = [_CHAT_TOOL, 5] if index % 2 else [_CHAT_TOOL] + return _Call( + "/v1/chat/completions", + { + "model": _OWNED_CHAT, + "messages": [{"role": "user", "content": marker}], + "stream": index % 4 == 1, + "tools": tools, + }, + marker, + "dropped" if index % 2 else "valid", + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_worker_sigkill_mid_burst_leaves_the_replacement_dropping_bad_entries(owned: _Owned) -> None: + proxy_url: Final = owned.proxy.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.proxy.log.read_text())), + lambda pids: len(pids) >= 2, + seconds=30, + ) + started_before: Final = len(workers) + serving: Final = owned.proxy.worker_pids() + owned.wire.drain() + with ExitStack() as stack: + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + ports: Final = tuple(_local_port(client) for client in clients) + calls: Final = tuple(_owned_call(index, "worker kill") for index in range(len(clients))) + owned.release.clear() + owned.holding.set() + stack.callback(owned.holding.clear) + stack.callback(owned.release.set) + futures: Final = tuple( + pool.submit(_status_or_dropped, client, owned.proxy.gateway.key, call) + for client, call in zip(clients, calls, strict=True) + ) + eventually(owned.held.qsize, lambda size: size == len(clients), seconds=60) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in serving} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in serving if shares[pid]), key=lambda pid: len(shares[pid])) + with owned.proxy.replacing_worker(victim) as kill: + kill() + owned.holding.clear() + owned.release.set() + results: Final = tuple(future.result(timeout=120) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), (port, result, shares) + eventually( + lambda: len(_STARTED_WORKER.findall(owned.proxy.log.read_text())), + lambda started: started > started_before, + _RESPAWN_SECONDS, + ) + second_wave: Final = tuple(_owned_call(index * 2 + 1, "after kill") for index in range(6)) + (fresh,) = _clients(stack, str(proxy_url), 1) + assert tuple(_post(fresh, owned.proxy.gateway.key, call).status_code for call in second_wave) == (200,) * len( + second_wave + ) + deliveries: Final = _deliveries(owned.wire) + everything: Final = calls + second_wave + assert tuple(_seen(deliveries, call.marker) for call in everything) == (1,) * len(everything) + assert tuple(_delivered_tool_names(deliveries, call.marker) for call in everything) == ( + ("get_weather",), + ) * len(everything) + assert owned.proxy.process.poll() is None diff --git a/tests/integration/providers/test_malformed_tools_wire.py b/tests/integration/providers/test_malformed_tools_wire.py new file mode 100644 index 00000000000..25a86f7fc7c --- /dev/null +++ b/tests/integration/providers/test_malformed_tools_wire.py @@ -0,0 +1,2731 @@ +from __future__ import annotations + +import asyncio +import http.client +import json +import os +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, replace +from itertools import chain, product +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from google import genai +from google.genai import types +from integration._support.anthropic_sse import error_body, message_json, message_stream +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from integration.authorization._guardrail_opt_out import upstream_observations +from pydantic import JsonValue, TypeAdapter + +from tests.integration.caching.response_cache_case import await_entry, redis_store + +_OPENAI_BACKEND: Final = "gpt-4o-mini" +_CLAUDE_BACKEND: Final = "claude-haiku-4-5" +_GEMINI_BACKEND: Final = "gemini-2.5-flash" +_PROVIDER_KEY: Final = "scripted-malformed-tools-key" +_COUNT: Final = 37 +_CONFIG_STORE: Final = "vs_integration_config_store" +_OTHER_STORE: Final = "vs_malformed_tools_other_store" +_CLAUDE_CODE_AGENT: Final = "claude-cli/2.1.0 (external, cli)" +_CODEX_AGENT: Final = "codex_cli_rs/0.50.0 (Linux 6.8; x86_64) xterm-256color" +_REJECT_MARK: Final = "upstream-401" +_DOWN: Final = "/down/" +_REFUSALS: Final = MappingProxyType( + { + "context_window_fallbacks": ( + "/context-full/", + { + "message": "This model's maximum context length is 8192 tokens.", + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded", + }, + ), + "content_policy_fallbacks": ( + "/content-policy/", + { + "message": "Your request was rejected as a result of our safety system.", + "type": "invalid_request_error", + "param": None, + "code": "content_policy_violation", + }, + ), + } +) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SPEND_ROWS: Final = ( + 'SELECT call_type, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id = %s OR litellm_call_id = %s' +) +_NOT_A_LIST: Final = "tools: Input should be a valid list" +_NOT_AN_OBJECT: Final = "tools.0: Input should be a valid dictionary" +_SECOND_NOT_AN_OBJECT: Final = "tools.1: Input should be a valid dictionary" +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_TOOL_FIELDS: Final = ("tools", "tool_choice", "parallel_tool_calls", "toolConfig", "tool_config") +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} + + +def _marker() -> str: + return f"marker-{uuid.uuid4().hex}" + + +def _tag(text: str) -> str: + found: Final = _MARKER.search(text) + return found.group(1) if found else "none" + + +def _answer(tag: str) -> str: + return f"answer {tag}" + + +def _gemini_body(tag: str) -> dict[str, JsonValue]: + return { + "responseId": f"gemini-{tag}-{uuid.uuid4().hex}", + "candidates": [ + {"content": {"role": "model", "parts": [{"text": _answer(tag)}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11}, + "modelVersion": _GEMINI_BACKEND, + } + + +def _interaction_body(tag: str) -> dict[str, JsonValue]: + return { + "id": f"interaction-{tag}", + "object": "interaction", + "model": _GEMINI_BACKEND, + "status": "completed", + "created": "2026-10-15T00:00:00Z", + "steps": [{"type": "content", "content": {"type": "text", "text": _answer(tag)}}], + "usage": {"total_tokens": 11}, + } + + +def _rejected_by_provider(path: str) -> Reply: + if path.startswith("/v1/messages"): + return Reply(status=401, body=error_body(401, "scripted 401")) + if path == "/v1beta/interactions": + return Reply( + status=401, + body=json.dumps({"error": {"code": 401, "message": "scripted 401", "status": "UNAUTHENTICATED"}}).encode(), + ) + return openai_error(401) + + +def _unavailable(path: str) -> Reply: + if "/v1/messages" in path: + return Reply(status=503, body=error_body(503, "scripted outage")) + return openai_error(503) + + +def _respond(request: Request) -> Reply: + text: Final = request.body.decode() + tag: Final = _tag(text) + path: Final = request.target.split("?")[0] + if path.startswith(_DOWN): + return _unavailable(path) + refusal: Final = next((error for prefix, error in _REFUSALS.values() if path.startswith(prefix)), None) + if refusal is not None: + return Reply(status=400, body=json.dumps({"error": refusal}).encode()) + if _REJECT_MARK in text: + return _rejected_by_provider(path) + stream: Final = bool(text) and _JSON_OBJECT.validate_json(request.body).get("stream") is True + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/responses": + return responses_reply(f"resp_{tag}", _OPENAI_BACKEND, _answer(tag), stream=stream) + if path == "/v1/messages/count_tokens": + return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) + if path == "/v1/messages" and stream: + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag)) + ) + if path == "/v1/messages": + return Reply(body=message_json(f"msg_{tag}", _CLAUDE_BACKEND, _answer(tag))) + if path.endswith(":streamGenerateContent"): + return Reply( + content_type="text/event-stream", chunks=(f"data: {json.dumps(_gemini_body(tag))}\r\n\r\n".encode(),) + ) + if path.endswith(":generateContent"): + return Reply(body=json.dumps(_gemini_body(tag)).encode()) + if path == "/v1beta/interactions": + return Reply(body=json.dumps(_interaction_body(tag)).encode()) + return Reply(status=404, body=json.dumps({"error": f"unscripted {request.target}"}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + wire: Wire + openai: str + anthropic: str + gemini: str + drop_group: str + gemini_drop_group: str + mixed_group: str + field_group: str + tools_drop_group: str + + +def _deployment_ids(scenario: Scenario, name: str, deployments: tuple[dict[str, JsonValue], ...]) -> tuple[str, ...]: + created: Final = tuple( + scenario.gateway.post("/model/new", {"model_name": name, "litellm_params": params, "model_info": {}}) + for params in deployments + ) + identities: Final = tuple(string_value(object_value(deployment["model_info"])["id"]) for deployment in created) + for identity in identities: + scenario.cleanups.callback(scenario.delete_model, identity) + return identities + + +def _same_name_deployments(scenario: Scenario, api_base: str, variants: tuple[Mapping[str, JsonValue], ...]) -> str: + name: Final = f"integration-mixed-{uuid.uuid4().hex}" + _deployment_ids( + scenario, + name, + tuple( + {"model": f"openai/{_OPENAI_BACKEND}", "api_key": _PROVIDER_KEY, "api_base": api_base, **variant} + for variant in variants + ), + ) + return name + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + with gateway_from_environment() as gateway, wire_server(_respond) as wire, gateway.scenario() as scenario: + openai_base: Final = f"{wire.url}/v1" + built: Final = _Rig( + gateway=gateway, + wire=wire, + openai=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY), + anthropic=scenario.model(model=f"anthropic/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + gemini=scenario.model(model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY), + drop_group=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY, drop_params=True), + gemini_drop_group=scenario.model( + model=f"gemini/{_GEMINI_BACKEND}", api_base=wire.url, api_key=_PROVIDER_KEY, drop_params=True + ), + mixed_group=_same_name_deployments(scenario, openai_base, ({"drop_params": True}, {})), + field_group=scenario.model(api_base=openai_base, api_key=_PROVIDER_KEY, additional_drop_params=["tools"]), + tools_drop_group=scenario.model( + api_base=openai_base, api_key=_PROVIDER_KEY, drop_params=True, tools=[CHAT.tool] + ), + ) + for group, deployments in ( + (built.openai, 1), + (built.anthropic, 1), + (built.gemini, 1), + (built.drop_group, 1), + (built.gemini_drop_group, 1), + (built.mixed_group, 2), + (built.field_group, 1), + (built.tools_drop_group, 1), + ): + wait_for_worker_readiness(gateway, group, deployments, _PROXY_WORKERS) + yield built + + +@pytest.fixture(autouse=True) +def _drained(rig: _Rig) -> None: + rig.wire.drain() + + +@dataclass(frozen=True, slots=True) +class _FieldGroups: + anthropic: str + gemini: str + + +@pytest.fixture(scope="module") +def field_groups(rig: _Rig) -> Iterator[_FieldGroups]: + with rig.gateway.scenario() as scenario: + built: Final = _FieldGroups( + anthropic=scenario.model( + model=f"anthropic/{_CLAUDE_BACKEND}", + api_base=rig.wire.url, + api_key=_PROVIDER_KEY, + additional_drop_params=["tools"], + ), + gemini=scenario.model( + model=f"gemini/{_GEMINI_BACKEND}", + api_base=rig.wire.url, + api_key=_PROVIDER_KEY, + additional_drop_params=["tools"], + ), + ) + for group in (built.anthropic, built.gemini): + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + yield built + + +@pytest.fixture(scope="module") +def limited_key(rig: _Rig) -> Iterator[str]: + with rig.gateway.scenario() as scenario: + yield scenario.key(rpm_limit=0) + + +@dataclass(frozen=True, slots=True) +class _Family: + name: str + tool: dict[str, JsonValue] + model: Callable[[_Rig], str] + path: Callable[[str, bool], str] + body: Callable[[str, str, bool], dict[str, JsonValue]] + target: Callable[[bool], str] + call_type: Callable[[bool], str] + upstream_id: Callable[[str], str | None] + tool_choice: tuple[str, JsonValue] + exact_tools: bool + + +def _chat_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], "stream": stream} + + +def _messages_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream} + + +def _responses_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "input": marker, "stream": stream} + + +def _generate_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + + +def _generate_path(model: str, stream: bool) -> str: + return ( + f"/v1beta/models/{model}:streamGenerateContent?alt=sse" if stream else f"/v1beta/models/{model}:generateContent" + ) + + +CHAT: Final = _Family( + name="chat", + tool={ + "type": "function", + "function": {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA}, + }, + model=lambda rig: rig.openai, + path=lambda model, stream: "/v1/chat/completions", + body=_chat_body, + target=lambda stream: "/v1/chat/completions", + call_type=lambda stream: "acompletion", + upstream_id=lambda tag: f"chatcmpl-{tag}", + tool_choice=("tool_choice", {"type": "function", "function": {"name": "get_weather"}}), + exact_tools=True, +) +MESSAGES: Final = _Family( + name="messages", + tool={"name": "get_weather", "description": "Look up the weather", "input_schema": _WEATHER_SCHEMA}, + model=lambda rig: rig.anthropic, + path=lambda model, stream: "/v1/messages", + body=_messages_body, + target=lambda stream: "/v1/messages", + call_type=lambda stream: "anthropic_messages", + upstream_id=lambda tag: f"msg_{tag}", + tool_choice=("tool_choice", {"type": "tool", "name": "get_weather"}), + exact_tools=True, +) +RESPONSES: Final = _Family( + name="responses", + tool={ + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + }, + model=lambda rig: rig.openai, + path=lambda model, stream: "/v1/responses", + body=_responses_body, + target=lambda stream: "/v1/responses", + call_type=lambda stream: "aresponses", + upstream_id=lambda tag: None, + tool_choice=("tool_choice", {"type": "function", "name": "get_weather"}), + exact_tools=True, +) +GENERATE: Final = _Family( + name="generate_content", + tool={ + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up the weather", "parameters": _WEATHER_SCHEMA} + ] + }, + model=lambda rig: rig.gemini, + path=_generate_path, + body=_generate_body, + target=lambda stream: ":streamGenerateContent" if stream else ":generateContent", + call_type=lambda stream: "agenerate_content_stream" if stream else "agenerate_content", + upstream_id=lambda tag: None, + tool_choice=("toolConfig", {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["get_weather"]}}), + exact_tools=False, +) +_FAMILIES: Final = (CHAT, MESSAGES, RESPONSES, GENERATE) +_FAMILY_IDS: Final = tuple(family.name for family in _FAMILIES) +_STREAMS: Final = pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +_EACH_FAMILY: Final = pytest.mark.parametrize("family", _FAMILIES, ids=_FAMILY_IDS) + + +@dataclass(frozen=True, slots=True) +class _Route: + path: Callable[[str], str] + body: Callable[[str, str], dict[str, JsonValue]] + model: Callable[[_Rig], str] + + +_ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( + { + "chat": _Route(lambda model: "/v1/chat/completions", lambda m, k: _chat_body(m, k, False), CHAT.model), + "chat_unversioned": _Route(lambda model: "/chat/completions", lambda m, k: _chat_body(m, k, False), CHAT.model), + "chat_azure_alias": _Route( + lambda model: f"/openai/deployments/{model}/chat/completions", + lambda m, k: _chat_body(m, k, False), + CHAT.model, + ), + "chat_stream": _Route(lambda model: "/v1/chat/completions", lambda m, k: _chat_body(m, k, True), CHAT.model), + "messages": _Route(lambda model: "/v1/messages", lambda m, k: _messages_body(m, k, False), MESSAGES.model), + "messages_stream": _Route( + lambda model: "/v1/messages", lambda m, k: _messages_body(m, k, True), MESSAGES.model + ), + "responses": _Route(lambda model: "/v1/responses", lambda m, k: _responses_body(m, k, False), RESPONSES.model), + "responses_stream": _Route( + lambda model: "/v1/responses", lambda m, k: _responses_body(m, k, True), RESPONSES.model + ), + "responses_compact": _Route( + lambda model: "/v1/responses/compact", lambda m, k: {"model": m, "input": k}, RESPONSES.model + ), + "generate_content": _Route( + lambda model: _generate_path(model, False), lambda m, k: _generate_body(m, k, False), GENERATE.model + ), + "stream_generate_content": _Route( + lambda model: _generate_path(model, True), lambda m, k: _generate_body(m, k, True), GENERATE.model + ), + "interactions": _Route( + lambda model: "/v1beta/interactions", lambda m, k: {"model": m, "input": k}, lambda rig: rig.gemini + ), + } +) +_MALFORMED_TOOLS: Final[Mapping[str, tuple[JsonValue, str]]] = MappingProxyType( + { + "int": (5, _NOT_A_LIST), + "empty_string": ("", _NOT_A_LIST), + "five_kb_string": ("x" * 5120, _NOT_A_LIST), + "list_of_names": (["get_weather"], _NOT_AN_OBJECT), + "empty_object": ({}, _NOT_A_LIST), + "false": (False, _NOT_A_LIST), + "zero": (0, _NOT_A_LIST), + } +) +_MALFORMED_ENTRIES: Final[Mapping[str, tuple[JsonValue, str]]] = MappingProxyType( + { + "int_entry": (5, _NOT_AN_OBJECT), + "string_function": ( + {"type": "function", "function": "get_weather"}, + "tools.0.function: Input should be a valid dictionary", + ), + "string_function_parameters": ( + {"type": "function", "function": {"name": "get_weather", "parameters": "bad"}}, + "tools.0.function.parameters: Input should be a valid dictionary", + ), + "list_input_schema": ( + {"name": "get_weather", "input_schema": [1]}, + "tools.0.input_schema: Input should be a valid dictionary", + ), + "int_parameters": ( + {"type": "function", "name": "get_weather", "parameters": 5}, + "tools.0.parameters: Input should be a valid dictionary", + ), + "string_vector_store_ids": ( + {"type": "file_search", "vector_store_ids": _OTHER_STORE}, + "tools.0.vector_store_ids: Input should be a valid list", + ), + } +) + + +def _send( + rig: _Rig, + family: _Family, + marker: str, + stream: bool, + *, + model: str | None = None, + key: str | None = None, + call_id: str | None = None, + user_agent: str | None = None, + **fields: JsonValue, +) -> httpx.Response: + chosen: Final = model or family.model(rig) + headers: Final = { + **({"x-litellm-call-id": call_id} if call_id else {}), + **({"user-agent": user_agent} if user_agent else {}), + } + return rig.gateway.request( + "POST", family.path(chosen, stream), {**family.body(chosen, marker, stream), **fields}, key=key, headers=headers + ) + + +def _delivered(rig: _Rig, marker: str) -> tuple[Request, ...]: + return tuple(request for request in rig.wire.drain() if marker in request.body.decode()) + + +def _one_delivery(rig: _Rig, marker: str, target: str) -> dict[str, JsonValue]: + (sent,) = _delivered(rig, marker) + assert sent.method == "POST" and target in sent.target, sent.target + return _JSON_OBJECT.validate_json(sent.body) + + +def _tool_names(family: _Family, body: Mapping[str, JsonValue]) -> tuple[str, ...]: + tools: Final = body.get("tools") + if not isinstance(tools, list): + return () + if family is GENERATE: + declarations: Final = chain.from_iterable(_declarations(object_value(tool)) for tool in tools) + return tuple(string_value(object_value(declaration)["name"]) for declaration in declarations) + if family is CHAT: + return tuple(string_value(object_value(object_value(tool)["function"])["name"]) for tool in tools) + return tuple(string_value(object_value(tool)["name"]) for tool in tools) + + +def _declarations(tool: Mapping[str, JsonValue]) -> list[JsonValue]: + declarations: Final = tool.get("functionDeclarations", tool.get("function_declarations")) + assert isinstance(declarations, list), tool + return declarations + + +def _assert_answered(family: _Family, response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + tag: Final = _tag(marker) + assert _answer(tag) in response.text, response.text + upstream_id: Final = family.upstream_id(tag) + if upstream_id is not None: + assert upstream_id in response.text, response.text + if family is RESPONSES and not response.headers.get("content-type", "").startswith("text/event-stream"): + assert same_response(string_value(_JSON_OBJECT.validate_json(response.content)["id"]), f"resp_{tag}") + + +def _rejection(response: httpx.Response) -> tuple[str, str, JsonValue]: + body: Final = _JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + if body.get("type") == "error": + return string_value(error["type"]), string_value(error["message"]), None + return string_value(error["type"]), string_value(error["message"]), error.get("param") + + +def _assert_rejected(response: httpx.Response, message: str, param: str = "tools") -> None: + assert response.status_code == 400, response.text + kind, said, named = _rejection(response) + assert kind == "invalid_request_error", response.text + assert message in said, response.text + assert named in (param, None), response.text + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows(_SPEND_ROWS, (call_id, call_id)) + + +def _logged_once(call_id: str) -> dict[str, JsonValue]: + (row,) = held(lambda: _spend_rows(call_id), lambda rows: len(rows) == 1, holding=3, seconds=70) + return row + + +def _tool_values(sent: Mapping[str, JsonValue]) -> frozenset[str]: + return frozenset(name for name in _TOOL_FIELDS if sent.get(name) is not None) + + +def _call_id() -> str: + return f"malformed-tools-{uuid.uuid4().hex}" + + +@_STREAMS +@_EACH_FAMILY +def test_well_formed_tools_reach_the_provider_and_log_one_success_row(rig: _Rig, family: _Family, stream: bool) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, stream, call_id=call_id, tools=[family.tool]) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(stream)) + assert _tool_names(family, sent) == ("get_weather",), sent + if family.exact_tools: + assert sent["tools"] == [family.tool], sent + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(stream), "success"), row + + +def test_openai_sdk_chat_with_tools_answers_with_the_provider_reply(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + completion: Final = client.chat.completions.create( + model=rig.openai, + messages=[{"role": "user", "content": marker}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + }, + } + ], + ) + assert completion.id == f"chatcmpl-{_tag(marker)}", completion + assert completion.choices[0].message.content == _answer(_tag(marker)), completion + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +def test_openai_sdk_rejects_malformed_tools_with_a_400(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with pytest.raises(openai.BadRequestError) as raised: + client.chat.completions.create( + model=rig.openai, messages=[{"role": "user", "content": marker}], extra_body={"tools": 5} + ) + assert raised.value.status_code == 400, raised.value + assert object_value(raised.value.body)["param"] == "tools", raised.value.body + assert _delivered(rig, marker) == () + + +def test_async_openai_sdk_streams_responses_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + + async def stream_text() -> str: + client: Final = openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False), + ) + events: Final = await client.responses.create( + model=rig.openai, + input=marker, + stream=True, + tools=[ + { + "type": "function", + "name": "get_weather", + "description": "Look up the weather", + "parameters": _WEATHER_SCHEMA, + "strict": False, + } + ], + ) + completed: Final = [event async for event in events if event.type == "response.completed"] + assert len(completed) == 1, completed + return completed[0].response.output_text + + assert asyncio.run(stream_text()) == _answer(_tag(marker)) + assert _tool_names(RESPONSES, _one_delivery(rig, marker, "/v1/responses")) == ("get_weather",) + + +def test_anthropic_sdk_streams_messages_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with client.messages.stream( + model=rig.anthropic, + max_tokens=16, + messages=[{"role": "user", "content": marker}], + tools=[{"name": "get_weather", "description": "Look up the weather", "input_schema": _WEATHER_SCHEMA}], + ) as stream: + final: Final = stream.get_final_message() + assert final.id == f"msg_{_tag(marker)}", final + assert _one_delivery(rig, marker, "/v1/messages")["tools"] == [MESSAGES.tool] + + +def test_google_genai_sdk_generates_with_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = genai.Client(api_key=rig.gateway.key, http_options={"base_url": str(rig.gateway.client.base_url)}) + response: Final = client.models.generate_content( + model=rig.gemini, + contents=marker, + config=types.GenerateContentConfig( + tools=[ + types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="get_weather", + description="Look up the weather", + parameters_json_schema=_WEATHER_SCHEMA, + ) + ] + ) + ], + automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True), + ), + ) + assert response.text == _answer(_tag(marker)), response + assert _tool_names(GENERATE, _one_delivery(rig, marker, ":generateContent")) == ("get_weather",) + + +@pytest.mark.parametrize("value", _MALFORMED_TOOLS.values(), ids=_MALFORMED_TOOLS.keys()) +@pytest.mark.parametrize("route", _ROUTES.values(), ids=_ROUTES.keys()) +def test_malformed_tools_answer_400_and_never_reach_the_provider( + rig: _Rig, route: _Route, value: tuple[JsonValue, str] +) -> None: + marker: Final = _marker() + model: Final = route.model(rig) + tools, message = value + response: Final = rig.gateway.request("POST", route.path(model), {**route.body(model, marker), "tools": tools}) + _assert_rejected(response, message) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("entry", _MALFORMED_ENTRIES.values(), ids=_MALFORMED_ENTRIES.keys()) +@pytest.mark.parametrize("family", (CHAT, MESSAGES), ids=("chat", "messages")) +def test_a_malformed_tool_entry_answers_400_naming_its_index( + rig: _Rig, family: _Family, entry: tuple[JsonValue, str] +) -> None: + marker: Final = _marker() + tool, message = entry + _assert_rejected(_send(rig, family, marker, False, tools=[tool]), message) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("value", [5, "x", []], ids=["int", "string", "list"]) +def test_malformed_web_search_options_answer_400(rig: _Rig, value: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, web_search_options=value) + expected: Final = "web_search_options: Input should be a valid dictionary" + _assert_rejected(response, expected, param="web_search_options") + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("value", [5, "x", []], ids=["int", "string", "list"]) +def test_malformed_web_search_options_drop_under_drop_params(rig: _Rig, value: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, web_search_options=value, drop_params=True) + _assert_answered(CHAT, response, marker) + assert "web_search_options" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +@_EACH_FAMILY +def test_wrong_key_with_malformed_tools_answers_401(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, False, key=f"sk-wrong-{uuid.uuid4().hex}", tools=5) + assert response.status_code == 401, response.text + assert _delivered(rig, marker) == () + + +def _count_body(rig: _Rig, marker: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"model": rig.anthropic, "messages": [{"role": "user", "content": marker}], **fields} + + +@pytest.mark.parametrize( + "fields", + [ + {"tools": 5}, + {"tools": ""}, + {"tools": "x" * 5120}, + {"tools": ["get_weather"]}, + {"messages": "x"}, + ], + ids=["int_tools", "empty_string_tools", "five_kb_tools", "list_of_names", "string_messages"], +) +def test_count_tokens_answers_an_anthropic_400_and_keeps_the_connection( + rig: _Rig, fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + port: Final = rig.gateway.client.base_url.port + assert port is not None + connection: Final = http.client.HTTPConnection(rig.gateway.client.base_url.host, port, timeout=30) + headers: Final = {"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"} + try: + connection.request( + "POST", "/v1/messages/count_tokens", body=json.dumps(_count_body(rig, marker, **fields)), headers=headers + ) + counted: Final = connection.getresponse() + counted_body: Final = _JSON_OBJECT.validate_json(counted.read()) + socket_after_count: Final = connection.sock + connection.request( + "POST", "/v1/messages", body=json.dumps(_messages_body(rig.anthropic, marker, False)), headers=headers + ) + followed: Final = connection.getresponse() + followed_body: Final = followed.read().decode() + socket_after_follow_up: Final = connection.sock + finally: + connection.close() + assert counted.status == 400, counted_body + assert counted_body["type"] == "error", counted_body + error: Final = object_value(counted_body["error"]) + assert error["type"] == "invalid_request_error", counted_body + assert string_value(error["message"]).startswith(next(iter(fields))), counted_body + assert followed.status == 200, followed_body + assert f"msg_{_tag(marker)}" in followed_body, followed_body + assert socket_after_count is not None and socket_after_follow_up is socket_after_count + (sent,) = _delivered(rig, marker) + assert sent.target == "/v1/messages", sent.target + + +def test_anthropic_sdk_count_tokens_rejects_malformed_tools(rig: _Rig) -> None: + marker: Final = _marker() + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.count_tokens( + model=rig.anthropic, messages=[{"role": "user", "content": marker}], extra_body={"tools": 5} + ) + assert raised.value.status_code == 400, raised.value + assert object_value(object_value(raised.value.body)["error"])["message"] == _NOT_A_LIST, raised.value.body + assert _delivered(rig, marker) == () + + +def test_count_tokens_forwards_well_formed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/messages/count_tokens", _count_body(rig, marker, tools=[MESSAGES.tool]) + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +@pytest.mark.parametrize( + ("headers", "fields"), + [({}, {"drop_params": True}), ({"user-agent": _CLAUDE_CODE_AGENT}, {})], + ids=["drop_params", "claude_code_agent"], +) +def test_count_tokens_drops_only_the_bad_entry( + rig: _Rig, headers: dict[str, str], fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + _count_body(rig, marker, tools=[MESSAGES.tool, 5], **fields), + headers=headers, + ) + assert response.status_code == 200, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +def test_responses_input_tokens_counts_well_formed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/input_tokens", {"model": rig.anthropic, "input": marker, "tools": [RESPONSES.tool]} + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + }, response.text + sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens") + assert "get_weather" in json.dumps(sent["tools"]), sent + + +def test_responses_input_tokens_rejects_malformed_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/input_tokens", {"model": rig.anthropic, "input": marker, "tools": 5} + ) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +def test_responses_input_tokens_drops_only_the_bad_entry(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": rig.anthropic, "input": marker, "tools": [RESPONSES.tool, 5], "drop_params": True}, + ) + assert response.status_code == 200, response.text + sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens") + tools: Final = sent["tools"] + assert isinstance(tools, list) and len(tools) == 1 and "get_weather" in json.dumps(tools), sent + + +def _allowlisted_key(scenario: Scenario) -> str: + return scenario.key(metadata={"allowed_tools": ["get_weather"]}) + + +@_EACH_FAMILY +def test_allowlisted_key_keeps_its_allowlist_with_malformed_tools(rig: _Rig, family: _Family) -> None: + with rig.gateway.scenario() as scenario: + key: Final = _allowlisted_key(scenario) + allowed: Final = _marker() + _assert_answered(family, _send(rig, family, allowed, False, key=key, tools=[family.tool]), allowed) + assert len(_delivered(rig, allowed)) == 1 + malformed: Final = _marker() + _assert_rejected(_send(rig, family, malformed, False, key=key, tools=5), _NOT_A_LIST) + denied_tool: Final = json.loads(json.dumps(family.tool).replace("get_weather", "delete_files")) + denied: Final = _marker() + refused: Final = _send(rig, family, denied, False, key=key, tools=[denied_tool]) + assert refused.status_code == 403, refused.text + assert "delete_files" in refused.text, refused.text + assert _delivered(rig, malformed) + _delivered(rig, denied) == () + + +def _file_search(store_ids: JsonValue) -> dict[str, JsonValue]: + return {"type": "file_search", "vector_store_ids": store_ids} + + +def test_store_limited_key_searches_its_store_from_a_file_search_tool(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search([_CONFIG_STORE])]) + _assert_answered(CHAT, response, marker) + sent: Final = _one_delivery(rig, marker, "/v1/chat/completions") + assert "file_search" not in json.dumps(sent), sent + searches: Final = eventually( + lambda: tuple( + observation + for observation in upstream_observations(rig.gateway) + if observation["path"] == f"/vector_stores/{_CONFIG_STORE}/search" and marker in json.dumps(observation) + ), + lambda found: len(found) == 1, + ) + assert len(searches) == 1, searches + + +@pytest.mark.parametrize( + ("store_ids", "status", "named"), + [ + ([_OTHER_STORE], 401, _OTHER_STORE), + (_CONFIG_STORE, 400, "tools.0.vector_store_ids: Input should be a valid list"), + (_OTHER_STORE, 401, _OTHER_STORE), + ([5, _OTHER_STORE], 401, _OTHER_STORE), + ([5, _CONFIG_STORE], 400, "tools.0.vector_store_ids.0: Input should be a valid string"), + ], + ids=["other_store", "allowed_store_as_string", "other_store_as_string", "int_beside_other", "int_beside_allowed"], +) +def test_store_limited_key_reads_every_string_store_id( + rig: _Rig, store_ids: JsonValue, status: int, named: str +) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search(store_ids)]) + assert response.status_code == status, response.text + assert named in response.text, response.text + assert _delivered(rig, marker) == () + + +def test_store_limited_key_treats_null_store_ids_as_none(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(object_permission={"vector_stores": [_CONFIG_STORE]}) + response: Final = _send(rig, CHAT, marker, False, key=key, tools=[_file_search(None)]) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [_file_search(None)] + + +@pytest.mark.parametrize( + ("first", "last", "status"), + [("5", "WELL_FORMED", 200), ("WELL_FORMED", "5", 400)], + ids=["well_formed_last", "malformed_last"], +) +def test_duplicate_tools_keys_use_the_last_one(rig: _Rig, first: str, last: str, status: int) -> None: + marker: Final = _marker() + well_formed: Final = json.dumps([CHAT.tool]) + prefix: Final = json.dumps({"model": rig.openai, "messages": [{"role": "user", "content": marker}]})[:-1] + raw: Final = ( + f'{prefix}, "tools": {first.replace("WELL_FORMED", well_formed)}, ' + f'"tools": {last.replace("WELL_FORMED", well_formed)}}}' + ) + response: Final = rig.gateway.client.post( + "/v1/chat/completions", + content=raw.encode(), + headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == status, response.text + delivered: Final = _delivered(rig, marker) + if status == 400: + _assert_rejected(response, _NOT_A_LIST) + assert delivered == () + return + (sent,) = delivered + assert _JSON_OBJECT.validate_json(sent.body)["tools"] == [CHAT.tool] + + +@pytest.mark.parametrize( + ("family", "route", "expected"), + [ + (CHAT, None, "acompletion"), + (MESSAGES, None, "anthropic_messages"), + (RESPONSES, None, "aresponses"), + (GENERATE, None, "agenerate_content"), + (None, "interactions", "acreate_interaction"), + ], + ids=["chat", "messages", "responses", "generate_content", "interactions"], +) +def test_a_malformed_tools_rejection_logs_one_failure_row_with_its_call_type( + rig: _Rig, family: _Family | None, route: str | None, expected: str +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = ( + _send(rig, family, marker, False, call_id=call_id, tools=5) + if family is not None + else rig.gateway.request( + "POST", + "/v1beta/interactions", + {"model": rig.gemini, "input": [{"type": "text", "text": marker}], "tools": 5}, + headers={"x-litellm-call-id": call_id}, + ) + ) + _assert_rejected(response, _NOT_A_LIST) + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (expected, "failure"), row + + +@pytest.mark.parametrize( + ("path", "body", "expected"), + [ + ("/v1/chat/completions", "chat", "acompletion"), + ("/v1/messages", "messages", "anthropic_messages"), + ("/v1beta/interactions", "interactions", "acreate_interaction"), + ], + ids=["chat", "messages", "interactions_list_input"], +) +def test_a_provider_401_with_well_formed_tools_logs_its_endpoint_call_type( + rig: _Rig, path: str, body: str, expected: str +) -> None: + marker: Final = f"{_marker()} {_REJECT_MARK}" + call_id: Final = _call_id() + payloads: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {**_chat_body(rig.openai, marker, False), "tools": [CHAT.tool]}, + "messages": {**_messages_body(rig.anthropic, marker, False), "tools": [MESSAGES.tool]}, + "interactions": { + "model": rig.gemini, + "input": [{"type": "text", "text": marker}], + "tools": [{"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}], + }, + } + response: Final = rig.gateway.request("POST", path, payloads[body], headers={"x-litellm-call-id": call_id}) + assert response.status_code == 401, response.text + assert len(_delivered(rig, marker)) == 1 + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (expected, "failure"), row + + +@pytest.mark.parametrize("drop", [True, "true"], ids=["bool", "string"]) +@_EACH_FAMILY +def test_request_drop_params_drops_malformed_tools_and_their_companions( + rig: _Rig, family: _Family, drop: JsonValue +) -> None: + marker: Final = _marker() + companion, value = family.tool_choice + response: Final = _send(rig, family, marker, False, tools=5, drop_params=drop, **{companion: value}) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(False)) + assert _tool_values(sent) == frozenset(), sent + + +@_EACH_FAMILY +def test_drop_params_keeps_the_good_entries_and_the_tool_choice(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + companion, value = family.tool_choice + response: Final = _send(rig, family, marker, False, tools=[family.tool, 5], drop_params=True, **{companion: value}) + _assert_answered(family, response, marker) + sent: Final = _one_delivery(rig, marker, family.target(False)) + assert _tool_names(family, sent) == ("get_weather",), sent + if family.exact_tools: + assert sent["tools"] == [family.tool], sent + assert companion in sent, sent + + +def test_additional_drop_params_naming_tools_drops_them(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, tools=5, additional_drop_params=["tools"]) + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_additional_drop_params_naming_another_field_still_rejects_tools(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, tools=5, additional_drop_params=["web_search_options"]) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize("family", (CHAT, GENERATE), ids=("chat", "generate_content")) +def test_a_drop_params_deployment_group_drops_malformed_tools(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + model: Final = rig.drop_group if family is CHAT else rig.gemini_drop_group + response: Final = _send(rig, family, marker, False, model=model, tools=5) + _assert_answered(family, response, marker) + assert _tool_values(_one_delivery(rig, marker, family.target(False))) == frozenset() + + +def test_a_drop_params_deployment_keeps_its_own_tools_when_it_drops_malformed_request_tools(rig: _Rig) -> None: + marker: Final = _marker() + companion, value = CHAT.tool_choice + response: Final = _send( + rig, CHAT, marker, False, model=rig.tools_drop_group, tools=5, parallel_tool_calls=False, **{companion: value} + ) + _assert_answered(CHAT, response, marker) + sent: Final = _one_delivery(rig, marker, "/v1/chat/completions") + assert _tool_names(CHAT, sent) == _tool_names(CHAT, {"tools": [CHAT.tool]}), sent + assert _tool_values(sent) == frozenset({"tools"}), sent + + +@pytest.mark.parametrize("drop", [False, None], ids=["false", "null"]) +def test_an_explicit_request_drop_params_overrides_the_deployment_group(rig: _Rig, drop: JsonValue) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=rig.drop_group, tools=5, drop_params=drop) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + + +def test_a_mixed_group_answers_malformed_tools_by_the_drop_params_of_the_deployment_it_picks(rig: _Rig) -> None: + markers: Final = tuple(_marker() for _ in range(20)) + statuses: Final = tuple( + _send(rig, CHAT, marker, False, model=rig.mixed_group, tools=5).status_code for marker in markers + ) + sent: Final = tuple(request.body.decode() for request in rig.wire.drain()) + deliveries: Final = tuple( + tuple(_JSON_OBJECT.validate_json(body) for body in sent if marker in body) for marker in markers + ) + assert sorted(set(statuses)) == [200, 400], statuses + assert all((status, len(bodies)) in ((200, 1), (400, 0)) for status, bodies in zip(statuses, deliveries)), ( + statuses, + deliveries, + ) + assert all(all("tools" not in body for body in bodies) for bodies in deliveries), deliveries + + +def test_a_deployment_dropping_tools_drops_every_tool(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool, 5]) + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_a_request_additional_drop_params_list_keeps_tools_a_deployment_list_drops(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=rig.field_group, + tools=[CHAT.tool], + additional_drop_params=["web_search_options"], + ) + _assert_answered(CHAT, response, marker) + assert _tool_names(CHAT, _one_delivery(rig, marker, "/v1/chat/completions")) == ("get_weather",) + + +def test_a_request_additional_drop_params_list_replaces_the_deployment_list(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=rig.field_group, + tools=[CHAT.tool, 5], + additional_drop_params=["web_search_options"], + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + ("family", "agent"), + [(MESSAGES, _CLAUDE_CODE_AGENT), (RESPONSES, _CODEX_AGENT), (CHAT, _CLAUDE_CODE_AGENT)], + ids=["claude_code_messages", "codex_responses", "claude_code_chat"], +) +def test_an_agentic_cli_drops_only_the_bad_entry(rig: _Rig, family: _Family, agent: str) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, False, user_agent=agent, tools=[family.tool, 5]) + _assert_answered(family, response, marker) + assert _one_delivery(rig, marker, family.target(False))["tools"] == [family.tool] + + +def test_an_agentic_cli_with_explicit_drop_params_false_rejects(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, MESSAGES, marker, False, user_agent=_CLAUDE_CODE_AGENT, tools=[MESSAGES.tool, 5], drop_params=False + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + "fields", [{"tools": None}, {"tools": []}, {}], ids=["null_tools", "empty_list_tools", "missing_tools"] +) +def test_empty_null_and_missing_tools_still_answer(rig: _Rig, fields: dict[str, JsonValue]) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, **fields) + _assert_answered(CHAT, response, marker) + assert not _one_delivery(rig, marker, "/v1/chat/completions").get("tools") + + +def test_null_web_search_options_answers_and_logs_success(rig: _Rig) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, CHAT, marker, False, call_id=call_id, web_search_options=None) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions").get("web_search_options") is None + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("acompletion", "success"), row + + +@_EACH_FAMILY +def test_request_drop_params_with_well_formed_tools_logs_one_success_row(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, False, call_id=call_id, tools=[family.tool], drop_params=True) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == ("get_weather",) + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(False), "success"), row + + +def test_repeated_identical_dropped_requests_each_log_once(rig: _Rig) -> None: + marker: Final = _marker() + call_ids: Final = tuple(_call_id() for _ in range(3)) + for call_id in call_ids: + response = _send(rig, GENERATE, marker, False, call_id=call_id, tools=[GENERATE.tool, 5], drop_params=True) + _assert_answered(GENERATE, response, marker) + delivered: Final = _delivered(rig, marker) + assert ( + tuple(_tool_names(GENERATE, _JSON_OBJECT.validate_json(sent.body)) for sent in delivered) + == (("get_weather",),) * 3 + ) + rows: Final = tuple(_logged_once(call_id) for call_id in call_ids) + assert tuple((row["call_type"], row["status"]) for row in rows) == (("agenerate_content", "success"),) * 3 + + +@pytest.mark.parametrize("family", (CHAT, MESSAGES, RESPONSES), ids=("chat", "messages", "responses")) +def test_a_cache_hit_twin_answers_from_the_cache_and_logs_once_each(rig: _Rig, family: _Family) -> None: + marker: Final = _marker() + first_id: Final = _call_id() + second_id: Final = _call_id() + first: Final = _send(rig, family, marker, False, call_id=first_id, tools=[family.tool]) + _assert_answered(family, first, marker) + await_entry(redis_store(), _answer(_tag(marker))) + second: Final = _send(rig, family, marker, False, call_id=second_id, tools=[family.tool]) + _assert_answered(family, second, marker) + assert len(_delivered(rig, marker)) == 1 + first_row: Final = _logged_once(first_id) + second_row: Final = _logged_once(second_id) + assert (first_row["call_type"], second_row["call_type"]) == (family.call_type(False),) * 2, (first_row, second_row) + assert second_row["cache_hit"] == "True", second_row + + +def test_a_dropped_request_shares_the_cache_entry_of_its_clean_twin(rig: _Rig) -> None: + marker: Final = _marker() + dropped: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool, 5], drop_params=True) + _assert_answered(CHAT, dropped, marker) + await_entry(redis_store(), _answer(_tag(marker))) + clean: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool], drop_params=True) + _assert_answered(CHAT, clean, marker) + (sent,) = _delivered(rig, marker) + assert _JSON_OBJECT.validate_json(sent.body)["tools"] == [CHAT.tool] + + +def _field_group(rig: _Rig, groups: _FieldGroups, family: _Family) -> str: + return {MESSAGES.name: groups.anthropic, GENERATE.name: groups.gemini}.get(family.name, rig.field_group) + + +_FIELD_GROUP_FAMILIES: Final = pytest.mark.parametrize( + "family", (MESSAGES, RESPONSES, GENERATE), ids=("messages", "responses", "generate_content") +) + + +@_FIELD_GROUP_FAMILIES +def test_each_family_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups, family: _Family +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + family, + marker, + False, + model=_field_group(rig, field_groups, family), + tools=[family.tool, 5], + additional_drop_params=["web_search_options"], + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + ("family", "forwarded"), + [(MESSAGES, ()), (RESPONSES, ()), (GENERATE, ("get_weather",))], + ids=("messages", "responses", "generate_content"), +) +def test_each_family_answers_malformed_tools_a_deployment_list_drops( + rig: _Rig, field_groups: _FieldGroups, family: _Family, forwarded: tuple[str, ...] +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, family, marker, False, model=_field_group(rig, field_groups, family), tools=[family.tool, 5] + ) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == forwarded + + +@_FIELD_GROUP_FAMILIES +def test_each_family_forwards_well_formed_tools_a_request_list_keeps( + rig: _Rig, field_groups: _FieldGroups, family: _Family +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + family, + marker, + False, + model=_field_group(rig, field_groups, family), + tools=[family.tool], + additional_drop_params=["web_search_options"], + ) + _assert_answered(family, response, marker) + assert _tool_names(family, _one_delivery(rig, marker, family.target(False))) == ("get_weather",) + + +_NOT_A_NAME_LIST: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "empty_list": [], + "null": None, + "string": "tools", + "empty_string": "", + "five_kb_string": "x" * 5120, + "int": 5, + } +) + + +@pytest.mark.parametrize("value", _NOT_A_NAME_LIST.values(), ids=_NOT_A_NAME_LIST.keys()) +def test_a_request_additional_drop_params_that_is_no_list_of_names_still_replaces_the_deployment_list( + rig: _Rig, value: JsonValue +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool, 5], additional_drop_params=value + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +@pytest.mark.parametrize( + "value", + [[], None, "tools", "", "x" * 5120], + ids=["empty_list", "null", "string", "empty_string", "five_kb_string"], +) +def test_a_request_additional_drop_params_replacing_the_deployment_list_keeps_well_formed_tools( + rig: _Rig, value: JsonValue +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=rig.field_group, tools=[CHAT.tool], additional_drop_params=value + ) + _assert_answered(CHAT, response, marker) + assert _one_delivery(rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +@pytest.mark.parametrize( + ("first", "last", "status"), + [('["tools"]', '["web_search_options"]', 400), ('["web_search_options"]', '["tools"]', 200)], + ids=["other_field_last", "tools_last"], +) +def test_duplicate_additional_drop_params_keys_use_the_last_one(rig: _Rig, first: str, last: str, status: int) -> None: + marker: Final = _marker() + prefix: Final = json.dumps( + {"model": rig.openai, "messages": [{"role": "user", "content": marker}], "tools": [CHAT.tool, 5]} + )[:-1] + response: Final = rig.gateway.client.post( + "/v1/chat/completions", + content=f'{prefix}, "additional_drop_params": {first}, "additional_drop_params": {last}}}'.encode(), + headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == status, response.text + if status == 400: + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + return + _assert_answered(CHAT, response, marker) + assert "tools" not in _one_delivery(rig, marker, "/v1/chat/completions") + + +def test_count_tokens_drops_the_bad_entry_a_deployment_list_names(rig: _Rig, field_groups: _FieldGroups) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {**_count_body(rig, marker, tools=[MESSAGES.tool, 5]), "model": field_groups.anthropic}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] + + +def test_count_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + { + **_count_body(rig, marker, tools=[MESSAGES.tool, 5], additional_drop_params=["web_search_options"]), + "model": field_groups.anthropic, + }, + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +def test_responses_input_tokens_drops_the_bad_entry_a_deployment_list_names( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": field_groups.anthropic, "input": marker, "tools": [RESPONSES.tool, 5]}, + ) + assert response.status_code == 200, response.text + tools: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] + assert isinstance(tools, list) and len(tools) == 1 and "get_weather" in json.dumps(tools), tools + + +def test_responses_input_tokens_rejects_malformed_tools_when_a_request_list_replaces_the_deployment_list( + rig: _Rig, field_groups: _FieldGroups +) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/input_tokens", + { + "model": field_groups.anthropic, + "input": marker, + "tools": [RESPONSES.tool, 5], + "additional_drop_params": ["web_search_options"], + }, + ) + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +def _text_items(marker: str) -> list[JsonValue]: + return [{"role": "user", "content": marker}] + + +@pytest.mark.parametrize("compact_input", [_text_items, lambda marker: marker], ids=["list_input", "string_input"]) +def test_a_compact_rejection_logs_one_failure_row_as_compact( + rig: _Rig, compact_input: Callable[[str], JsonValue] +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.request( + "POST", + "/v1/responses/compact", + {"model": rig.openai, "input": compact_input(marker), "tools": 5}, + headers={"x-litellm-call-id": call_id}, + ) + _assert_rejected(response, _NOT_A_LIST) + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("acompact_responses", "failure"), row + + +@dataclass(frozen=True, slots=True) +class _Limited: + path: Callable[[_Rig], str] + body: Callable[[_Rig, str], dict[str, JsonValue]] + call_type: str + + +_LIMITED: Final[Mapping[str, _Limited]] = MappingProxyType( + { + "chat": _Limited( + lambda rig: "/v1/chat/completions", lambda rig, k: _chat_body(rig.openai, k, False), "acompletion" + ), + "messages": _Limited( + lambda rig: "/v1/messages", lambda rig, k: _messages_body(rig.anthropic, k, False), "anthropic_messages" + ), + "responses_list_input": _Limited( + lambda rig: "/v1/responses", lambda rig, k: {"model": rig.openai, "input": _text_items(k)}, "aresponses" + ), + "compact_list_input": _Limited( + lambda rig: "/v1/responses/compact", + lambda rig, k: {"model": rig.openai, "input": _text_items(k)}, + "acompact_responses", + ), + "generate_content": _Limited( + lambda rig: _generate_path(rig.gemini, False), + lambda rig, k: _generate_body(rig.gemini, k, False), + "agenerate_content", + ), + "interactions_list_input": _Limited( + lambda rig: "/v1beta/interactions", + lambda rig, k: {"model": rig.gemini, "input": [{"type": "text", "text": k}]}, + "acreate_interaction", + ), + "video_generation": _Limited( + lambda rig: "/v1/videos", lambda rig, k: {"model": rig.openai, "prompt": k}, "avideo_generation" + ), + } +) + + +@pytest.mark.parametrize("limited", _LIMITED.values(), ids=_LIMITED.keys()) +def test_a_rate_limited_request_logs_one_failure_row_with_its_endpoint_call_type( + rig: _Rig, limited_key: str, limited: _Limited +) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.request( + "POST", + limited.path(rig), + limited.body(rig, marker), + key=limited_key, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 429, response.text + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (limited.call_type, "failure"), row + + +@pytest.mark.parametrize("family", (CHAT, MESSAGES), ids=("chat", "messages")) +def test_a_rate_limited_twin_of_a_cached_request_logs_its_endpoint_call_type( + rig: _Rig, limited_key: str, family: _Family +) -> None: + marker: Final = _marker() + _assert_answered(family, _send(rig, family, marker, False, tools=[family.tool]), marker) + await_entry(redis_store(), _answer(_tag(marker))) + call_id: Final = _call_id() + response: Final = _send(rig, family, marker, False, key=limited_key, call_id=call_id, tools=[family.tool]) + assert response.status_code == 429, response.text + assert len(_delivered(rig, marker)) == 1 + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == (family.call_type(False), "failure"), row + + +def test_a_rate_limited_image_edit_logs_one_failure_row_as_an_image_edit(rig: _Rig, limited_key: str) -> None: + marker: Final = _marker() + call_id: Final = _call_id() + response: Final = rig.gateway.client.post( + "/v1/images/edits", + data={"model": rig.openai, "prompt": marker}, + files={"image": ("edit.png", b"\x89PNG\r\n\x1a\n", "image/png")}, + headers={"Authorization": f"Bearer {limited_key}", "x-litellm-call-id": call_id}, + ) + assert response.status_code == 429, response.text + assert _delivered(rig, marker) == () + row: Final = _logged_once(call_id) + assert (row["call_type"], row["status"]) == ("aimage_edit", "failure"), row + + +_Sender = Callable[[_Rig, str, str, Mapping[str, JsonValue]], httpx.Response] + + +@dataclass(frozen=True, slots=True) +class _Checked: + provider: str + tool: dict[str, JsonValue] + named_as: _Family + send: _Sender + target: str + answered: Callable[[httpx.Response, str], None] + kept_by_a_tools_list: tuple[str, ...] + bad_entry: dict[str, JsonValue] + + +def _family_sender(family: _Family) -> _Sender: + return lambda rig, model, marker, fields: _send(rig, family, marker, False, model=model, **fields) + + +def _family_answered(family: _Family) -> Callable[[httpx.Response, str], None]: + return lambda response, marker: _assert_answered(family, response, marker) + + +def _count_tokens(rig: _Rig, model: str, marker: str, fields: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": marker}], **fields}, + ) + + +def _input_tokens(rig: _Rig, model: str, marker: str, fields: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": marker, **fields}) + + +def _counted(expected: Mapping[str, JsonValue]) -> Callable[[httpx.Response, str], None]: + def check(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == expected, (marker, response.text) + + return check + + +_CHAT_BAD_ENTRY: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": "bad"}, +} +_MESSAGES_BAD_ENTRY: Final[dict[str, JsonValue]] = {"name": "get_weather", "input_schema": [1]} +_RESPONSES_BAD_ENTRY: Final[dict[str, JsonValue]] = {"type": "function", "name": "get_weather", "parameters": 5} +_CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType( + { + "chat": _Checked( + provider="openai", + tool=CHAT.tool, + named_as=CHAT, + send=_family_sender(CHAT), + target="/v1/chat/completions", + answered=_family_answered(CHAT), + kept_by_a_tools_list=(), + bad_entry=_CHAT_BAD_ENTRY, + ), + "messages": _Checked( + provider="anthropic", + tool=MESSAGES.tool, + named_as=MESSAGES, + send=_family_sender(MESSAGES), + target="/v1/messages", + answered=_family_answered(MESSAGES), + kept_by_a_tools_list=(), + bad_entry=_MESSAGES_BAD_ENTRY, + ), + "responses": _Checked( + provider="openai", + tool=RESPONSES.tool, + named_as=RESPONSES, + send=_family_sender(RESPONSES), + target="/v1/responses", + answered=_family_answered(RESPONSES), + kept_by_a_tools_list=(), + bad_entry=_RESPONSES_BAD_ENTRY, + ), + "generate_content": _Checked( + provider="gemini", + tool=GENERATE.tool, + named_as=GENERATE, + send=_family_sender(GENERATE), + target=":generateContent", + answered=_family_answered(GENERATE), + kept_by_a_tools_list=("get_weather",), + bad_entry={"parameters": 5}, + ), + "count_tokens": _Checked( + provider="anthropic", + tool=MESSAGES.tool, + named_as=MESSAGES, + send=_count_tokens, + target="/v1/messages/count_tokens", + answered=_counted({"input_tokens": _COUNT}), + kept_by_a_tools_list=("get_weather",), + bad_entry=_MESSAGES_BAD_ENTRY, + ), + "input_tokens": _Checked( + provider="anthropic", + tool=RESPONSES.tool, + named_as=MESSAGES, + send=_input_tokens, + target="/v1/messages/count_tokens", + answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT}), + kept_by_a_tools_list=("get_weather",), + bad_entry=_RESPONSES_BAD_ENTRY, + ), + } +) +_EACH_CHECKED: Final = pytest.mark.parametrize("checked", _CHECKED.values(), ids=_CHECKED.keys()) +_CHAT_AND_COUNT: Final = pytest.mark.parametrize( + "checked", (_CHECKED["chat"], _CHECKED["count_tokens"]), ids=("chat", "count_tokens") +) +_ENTRY_SHAPES: Final = pytest.mark.parametrize( + "bad_entry", + [lambda checked: 5, lambda checked: checked.bad_entry], + ids=["int_entry", "dict_entry"], +) +_BadEntry = Callable[[_Checked], JsonValue] + + +def _assert_forwarded( + rig: _Rig, checked: _Checked, response: httpx.Response, marker: str, names: tuple[str, ...] +) -> None: + checked.answered(response, marker) + assert _tool_names(checked.named_as, _one_delivery(rig, marker, checked.target)) == names + + +def _assert_held_back(rig: _Rig, response: httpx.Response, marker: str) -> None: + _assert_rejected(response, _SECOND_NOT_AN_OBJECT) + assert _delivered(rig, marker) == () + + +_PROVIDERS: Final[Mapping[str, Callable[[str], dict[str, JsonValue]]]] = MappingProxyType( + { + "openai": lambda wire: { + "model": f"openai/{_OPENAI_BACKEND}", + "api_base": f"{wire}/v1", + "api_key": _PROVIDER_KEY, + }, + "anthropic": lambda wire: {"model": f"anthropic/{_CLAUDE_BACKEND}", "api_base": wire, "api_key": _PROVIDER_KEY}, + "gemini": lambda wire: {"model": f"gemini/{_GEMINI_BACKEND}", "api_base": wire, "api_key": _PROVIDER_KEY}, + } +) +_DROP_VARIANTS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "plain": {}, + "lists_web_search": {"additional_drop_params": ["web_search_options"]}, + "lists_tools": {"additional_drop_params": ["tools"]}, + "keeps": {"drop_params": False}, + } +) + + +@dataclass(frozen=True, slots=True) +class _Owned: + rig: _Rig + suffix: str + + +def _owned_name(suffix: str, variant: str, provider: str) -> str: + return f"malformed-tools-{variant}-{provider}-{suffix}" + + +def _owned_id(suffix: str, variant: str, provider: str) -> str: + return f"{_owned_name(suffix, variant, provider)}-id" + + +@dataclass(frozen=True, slots=True) +class _OwnedSetup: + router_settings: Callable[[str], dict[str, JsonValue]] + general: Mapping[str, JsonValue] = MappingProxyType({}) + extra_models: Callable[[str, str], tuple[dict[str, JsonValue], ...]] = lambda wire, suffix: () + + +def _router_defaults(defaults: dict[str, JsonValue]) -> _OwnedSetup: + return _OwnedSetup(router_settings=lambda suffix: {"default_litellm_params": defaults}) + + +def _owned_config(path: Path, wire: str, suffix: str, setup: _OwnedSetup) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + *( + { + "model_name": _owned_name(suffix, variant, provider), + "litellm_params": {**deployment(wire), **extra}, + "model_info": {"id": _owned_id(suffix, variant, provider)}, + } + for (variant, extra), (provider, deployment) in product( + _DROP_VARIANTS.items(), _PROVIDERS.items() + ) + ), + *setup.extra_models(wire, suffix), + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + **setup.general, + }, + "router_settings": setup.router_settings(suffix), + } + ) + ) + return path + + +@contextmanager +def _owned_proxy(rig: _Rig, directory: Path, setup: _OwnedSetup) -> Iterator[_Owned]: + suffix: Final = uuid.uuid4().hex[:12] + config: Final = _owned_config(directory / "owned.yaml", rig.wire.url, suffix, setup) + with owned_proxy_process(rig.gateway, directory, {}, config=config, workers=2) as proxy: + yield _Owned(replace(rig, gateway=proxy.gateway), suffix) + + +@pytest.fixture(scope="module") +def tools_default(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-tools-default") + with _owned_proxy(rig, directory, _router_defaults({"additional_drop_params": ["tools"]})) as owned: + yield owned + + +@pytest.fixture(scope="module") +def drop_default(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-drop-default") + with _owned_proxy(rig, directory, _router_defaults({"drop_params": True, "additional_drop_params": None})) as owned: + yield owned + + +@pytest.fixture(scope="module") +def client_side(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-client-side") + setup: Final = replace( + _router_defaults({"additional_drop_params": "tools"}), + general=MappingProxyType({"allow_client_side_credentials": True}), + ) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_list_naming_tools_answers_malformed_tools( + tools_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_EACH_CHECKED +def test_a_request_list_replaces_a_router_default_list(tools_default: _Owned, checked: _Checked) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", checked.provider) + response: Final = checked.send( + tools_default.rig, + model, + marker, + {"tools": [checked.tool, 5], "additional_drop_params": ["web_search_options"]}, + ) + _assert_held_back(tools_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_list_replaces_a_deployment_list( + tools_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "lists_web_search", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +def test_a_router_default_list_replaces_the_list_of_a_deployment_named_by_id( + tools_default: _Owned, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + model: Final = _owned_id(tools_default.suffix, "lists_web_search", checked.provider) + response: Final = checked.send(tools_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(tools_default.rig, checked, response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("value", [None, [], 5], ids=["null", "empty_list", "int"]) +def test_a_request_list_naming_nothing_replaces_a_router_default_list(tools_default: _Owned, value: JsonValue) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", "openai") + response: Final = _send( + tools_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], additional_drop_params=value + ) + _assert_held_back(tools_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_naming_tools_drops_well_formed_tools(tools_default: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(tools_default.suffix, "plain", "openai") + response: Final = _send(tools_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool]) + _assert_forwarded(tools_default.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_router_default_drop_params_keeps_the_good_entries( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "plain", checked.provider) + response: Final = checked.send(drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(drop_default.rig, checked, response, marker, ("get_weather",)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("drop", [False, None], ids=["false", "null"]) +def test_a_request_drop_params_replaces_a_router_default_drop_params(drop_default: _Owned, drop: JsonValue) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "plain", "openai") + response: Final = _send(drop_default.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], drop_params=drop) + _assert_held_back(drop_default.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_CHAT_AND_COUNT +def test_a_router_default_drop_params_replaces_a_deployment_drop_params( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "keeps", checked.provider) + response: Final = checked.send(drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(drop_default.rig, checked, response, marker, ("get_weather",)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_ENTRY_SHAPES +@_CHAT_AND_COUNT +def test_a_null_router_default_list_leaves_the_deployment_list_in_force( + drop_default: _Owned, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + model: Final = _owned_name(drop_default.suffix, "lists_tools", checked.provider) + response: Final = checked.send( + drop_default.rig, model, marker, {"tools": [checked.tool, bad_entry(checked)], "drop_params": False} + ) + _assert_forwarded(drop_default.rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_that_is_a_string_still_rejects_a_bad_tool_entry(client_side: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(client_side.suffix, "plain", "openai") + response: Final = _send(client_side.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5]) + _assert_held_back(client_side.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_router_default_list_that_is_a_string_keeps_well_formed_tools(client_side: _Owned) -> None: + marker: Final = _marker() + model: Final = _owned_name(client_side.suffix, "plain", "openai") + response: Final = _send(client_side.rig, CHAT, marker, False, model=model, tools=[CHAT.tool]) + _assert_answered(CHAT, response, marker) + assert _one_delivery(client_side.rig, marker, "/v1/chat/completions")["tools"] == [CHAT.tool] + + +_USER_CONFIG_DROPS: Final[Mapping[str, tuple[dict[str, JsonValue], dict[str, JsonValue], tuple[str, ...]]]] = ( + MappingProxyType( + { + "deployment_list": ({"additional_drop_params": ["tools"]}, {}, ()), + "router_default_list": ({}, {"additional_drop_params": ["tools"]}, ()), + "router_default_drop_params": ({}, {"drop_params": True}, ("get_weather",)), + } + ) +) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize(("deployment", "defaults", "kept"), _USER_CONFIG_DROPS.values(), ids=_USER_CONFIG_DROPS.keys()) +def test_a_user_config_router_decides_which_malformed_tools_drop( + client_side: _Owned, + deployment: dict[str, JsonValue], + defaults: dict[str, JsonValue], + kept: tuple[str, ...], +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + name: Final = f"malformed-tools-user-config-{uuid.uuid4().hex}" + user_config: Final[dict[str, JsonValue]] = { + "model_list": [ + {"model_name": name, "litellm_params": {**_PROVIDERS["openai"](client_side.rig.wire.url), **deployment}} + ], + "default_litellm_params": defaults, + } + response: Final = checked.send( + client_side.rig, name, marker, {"tools": [checked.tool, checked.bad_entry], "user_config": user_config} + ) + _assert_forwarded(client_side.rig, checked, response, marker, kept) + + +@dataclass(frozen=True, slots=True) +class _ById: + lists_tools: Mapping[str, str] + drops: str + mixed_drops: str + mixed_keeps: str + shadowed: str + + +@pytest.fixture(scope="module") +def by_id(rig: _Rig) -> Iterator[_ById]: + openai_deployment: Final = _PROVIDERS["openai"](rig.wire.url) + lists_tools: Final[dict[str, JsonValue]] = {"additional_drop_params": ["tools"]} + names: Final = MappingProxyType({provider: f"integration-by-id-{uuid.uuid4().hex}" for provider in _PROVIDERS}) + drops_name: Final = f"integration-by-id-drops-{uuid.uuid4().hex}" + mixed: Final = f"integration-by-id-mixed-{uuid.uuid4().hex}" + shadow_name: Final = f"integration-by-id-shadow-{uuid.uuid4().hex}" + with rig.gateway.scenario() as scenario: + listed: Final = MappingProxyType( + { + provider: _deployment_ids(scenario, names[provider], ({**deployment(rig.wire.url), **lists_tools},))[0] + for provider, deployment in _PROVIDERS.items() + } + ) + (drops,) = _deployment_ids(scenario, drops_name, ({**openai_deployment, "drop_params": True},)) + mixed_drops, mixed_keeps = _deployment_ids( + scenario, mixed, ({**openai_deployment, **lists_tools}, openai_deployment) + ) + (shadowed,) = _deployment_ids(scenario, shadow_name, ({**openai_deployment, **lists_tools},)) + _deployment_ids(scenario, shadowed, (openai_deployment,)) + for group, deployments in ( + *((name, 1) for name in names.values()), + (drops_name, 1), + (mixed, 2), + (shadow_name, 1), + (shadowed, 1), + ): + wait_for_worker_readiness(rig.gateway, group, deployments, _PROXY_WORKERS) + yield _ById(listed, drops, mixed_drops, mixed_keeps, shadowed) + + +@_ENTRY_SHAPES +@_EACH_CHECKED +def test_a_deployment_id_answers_malformed_tools_its_deployment_list_drops( + rig: _Rig, by_id: _ById, checked: _Checked, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, by_id.lists_tools[checked.provider], marker, {"tools": [checked.tool, bad_entry(checked)]} + ) + _assert_forwarded(rig, checked, response, marker, checked.kept_by_a_tools_list) + + +@_EACH_CHECKED +def test_a_request_list_replaces_the_list_of_a_deployment_named_by_id( + rig: _Rig, by_id: _ById, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, + by_id.lists_tools[checked.provider], + marker, + {"tools": [checked.tool, 5], "additional_drop_params": ["web_search_options"]}, + ) + _assert_held_back(rig, response, marker) + + +@_ENTRY_SHAPES +def test_the_id_of_the_dropping_deployment_in_a_mixed_group_drops_malformed_tools( + rig: _Rig, by_id: _ById, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send(rig, by_id.mixed_drops, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ()) + + +def test_the_keeping_id_of_a_mixed_group_rejects_malformed_tools(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.mixed_keeps, tools=[CHAT.tool, 5]) + _assert_held_back(rig, response, marker) + + +@_ENTRY_SHAPES +def test_the_id_of_a_drop_params_deployment_keeps_the_good_entries( + rig: _Rig, by_id: _ById, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send(rig, by_id.drops, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ("get_weather",)) + + +def test_a_model_name_wins_over_an_equal_deployment_id(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.shadowed, tools=[CHAT.tool, 5]) + _assert_held_back(rig, response, marker) + + +def test_a_deployment_id_with_a_tools_list_drops_well_formed_tools(rig: _Rig, by_id: _ById) -> None: + marker: Final = _marker() + response: Final = _send(rig, CHAT, marker, False, model=by_id.lists_tools["openai"], tools=[CHAT.tool]) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ()) + + +def _failed_over(request: Request) -> bool: + return request.target.startswith((_DOWN, *(prefix for prefix, _ in _REFUSALS.values()))) + + +def _assert_failed_over( + rig: _Rig, checked: _Checked, response: httpx.Response, marker: str, names: tuple[str, ...] +) -> None: + checked.answered(response, marker) + delivered: Final = _delivered(rig, marker) + (served,) = tuple(request for request in delivered if not _failed_over(request)) + assert checked.target in served.target, served.target + assert _tool_names(checked.named_as, _JSON_OBJECT.validate_json(served.body)) == names + + +_FALLBACK_CHECKED: Final = pytest.mark.parametrize( + "checked", + tuple(_CHECKED[name] for name in ("chat", "messages", "responses")), + ids=("chat", "messages", "responses"), +) +_TOKEN_COUNTS: Final = pytest.mark.parametrize( + "checked", (_CHECKED["count_tokens"], _CHECKED["input_tokens"]), ids=("count_tokens", "input_tokens") +) + + +@dataclass(frozen=True, slots=True) +class _Reroute: + down: Mapping[str, str] + lists_tools: Mapping[str, str] + refusing: Mapping[str, str] + fallback_keys: Mapping[str, str] + team_key: str + team_models: Mapping[str, str] + + +@pytest.fixture(scope="module") +def reroute(rig: _Rig, field_groups: _FieldGroups) -> Iterator[_Reroute]: + lists_tools: Final = MappingProxyType({"openai": rig.field_group, "anthropic": field_groups.anthropic}) + with rig.gateway.scenario() as scenario: + down: Final = MappingProxyType( + {provider: scenario.model(**_PROVIDERS[provider](f"{rig.wire.url}/down")) for provider in lists_tools} + ) + refusing: Final = MappingProxyType( + { + chain: scenario.model(**_PROVIDERS["openai"](f"{rig.wire.url}{prefix.rstrip('/')}")) + for chain, (prefix, _) in _REFUSALS.items() + } + ) + router_settings: Final[dict[str, JsonValue]] = {"fallbacks": [{down["openai"]: [rig.field_group]}]} + public_team: Final = scenario.team() + team_models: Final = MappingProxyType( + { + provider: scenario.model( + **_PROVIDERS[provider](rig.wire.url), + additional_drop_params=["tools"], + model_info={"team_id": public_team}, + ) + for provider in lists_tools + } + ) + built: Final = _Reroute( + down=down, + lists_tools=lists_tools, + refusing=refusing, + fallback_keys=MappingProxyType( + { + "key": scenario.key(router_settings=router_settings), + "team": scenario.key(team_id=scenario.team(router_settings=router_settings)), + } + ), + team_key=scenario.key(team_id=public_team), + team_models=team_models, + ) + for group in (*down.values(), *refusing.values(), *team_models.values()): + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + yield built + + +@_FALLBACK_CHECKED +def test_a_request_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, + reroute.down[checked.provider], + marker, + {"tools": [checked.tool, checked.bad_entry], "fallbacks": [reroute.lists_tools[checked.provider]]}, + ) + _assert_failed_over(rig, checked, response, marker, ()) + + +def test_a_generic_request_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["chat"] + response: Final = checked.send( + rig, + reroute.down["openai"], + marker, + {"tools": [checked.tool, checked.bad_entry], "fallbacks": [{"*": [rig.field_group]}]}, + ) + _assert_failed_over(rig, checked, response, marker, ()) + + +@pytest.mark.parametrize( + "fallbacks", + [[], None, lambda rig: [{rig.drop_group: [rig.field_group]}]], + ids=["empty_list", "null", "chain_of_another_group"], +) +def test_a_request_fallback_that_cannot_apply_still_rejects_a_bad_tool_entry( + rig: _Rig, fallbacks: JsonValue | Callable[[_Rig], JsonValue] +) -> None: + marker: Final = _marker() + chains: Final = fallbacks(rig) if callable(fallbacks) else fallbacks + response: Final = _send(rig, CHAT, marker, False, tools=[CHAT.tool, 5], fallbacks=chains) + _assert_held_back(rig, response, marker) + + +@pytest.mark.parametrize("owner", ["key", "team"]) +def test_router_settings_fallbacks_of_a_key_or_team_answer_a_bad_tool_entry( + rig: _Rig, reroute: _Reroute, owner: str +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=reroute.down["openai"], + key=reroute.fallback_keys[owner], + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_failed_over(rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.parametrize("chain", tuple(_REFUSALS)) +def test_an_error_specific_fallback_never_sees_a_bad_tool_entry_the_first_deployment_rejects( + rig: _Rig, reroute: _Reroute, chain: str +) -> None: + marker: Final = _marker() + refusing: Final = reroute.refusing[chain] + response: Final = _send( + rig, + CHAT, + marker, + False, + model=refusing, + tools=[CHAT.tool, 5], + **{chain: [{refusing: [rig.field_group]}]}, + ) + _assert_held_back(rig, response, marker) + + +def test_a_comma_batch_lets_each_group_decide_on_a_bad_tool_entry(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, CHAT, marker, False, model=f"{rig.field_group},{rig.openai}", tools=[CHAT.tool, _CHAT_BAD_ENTRY] + ) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + delivered: Final = sorted( + (_JSON_OBJECT.validate_json(request.body).get("tools") for request in _delivered(rig, marker)), key=bool + ) + assert delivered == [None] + + +def test_a_team_public_model_name_answers_a_bad_tool_entry_its_deployment_list_drops( + rig: _Rig, reroute: _Reroute +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + model=reroute.team_models["openai"], + key=reroute.team_key, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ()) + + +def test_count_tokens_by_a_team_public_model_name_rejects_a_bad_tool_entry(rig: _Rig, reroute: _Reroute) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": reroute.team_models["anthropic"], + "messages": [{"role": "user", "content": marker}], + "tools": [MESSAGES.tool, 5], + }, + key=reroute.team_key, + ) + _assert_held_back(rig, response, marker) + + +def test_a_request_list_naming_web_search_options_on_a_fallback_chain_still_rejects_a_bad_tool_entry( + rig: _Rig, +) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + tools=[CHAT.tool, 5], + web_search_options=5, + additional_drop_params=["web_search_options"], + fallbacks=[{rig.openai: [rig.field_group]}], + ) + _assert_held_back(rig, response, marker) + + +def test_request_drop_params_on_a_fallback_chain_keeps_the_good_entries(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _send( + rig, + CHAT, + marker, + False, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + drop_params=True, + fallbacks=[{rig.openai: [rig.field_group]}], + ) + _assert_forwarded(rig, _CHECKED["chat"], response, marker, ("get_weather",)) + + +@_ENTRY_SHAPES +def test_a_request_list_with_a_non_string_entry_drops_tools_it_names_on_messages( + rig: _Rig, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["messages"] + response: Final = checked.send( + rig, + rig.anthropic, + marker, + {"tools": [checked.tool, bad_entry(checked)], "additional_drop_params": ["tools", 5]}, + ) + _assert_forwarded(rig, checked, response, marker, ()) + + +@_ENTRY_SHAPES +def test_a_deployment_list_with_a_non_string_entry_drops_tools_it_names_on_messages( + rig: _Rig, bad_entry: _BadEntry +) -> None: + marker: Final = _marker() + checked: Final = _CHECKED["messages"] + with rig.gateway.scenario() as scenario: + group: Final = scenario.model(**_PROVIDERS["anthropic"](rig.wire.url), additional_drop_params=["tools", 5]) + wait_for_worker_readiness(rig.gateway, group, 1, _PROXY_WORKERS) + response: Final = checked.send(rig, group, marker, {"tools": [checked.tool, bad_entry(checked)]}) + _assert_forwarded(rig, checked, response, marker, ()) + + +@_TOKEN_COUNTS +def test_token_counting_with_a_request_list_with_a_non_string_entry_drops_the_bad_entry( + rig: _Rig, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send( + rig, rig.anthropic, marker, {"tools": [checked.tool, 5], "additional_drop_params": ["tools", 5]} + ) + _assert_forwarded(rig, checked, response, marker, ("get_weather",)) + + +def _alias(suffix: str, variant: str) -> str: + return f"malformed-tools-alias-{variant}-{suffix}" + + +def _routing_group(suffix: str, provider: str) -> str: + return f"malformed-tools-routed-{provider}-{suffix}" + + +def _wildcard(suffix: str) -> str: + return f"malformed-tools-wild-{suffix}" + + +def _direct(suffix: str) -> str: + return f"malformed-tools-direct-{suffix}" + + +def _down_models(wire: str, suffix: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + {"model_name": _owned_name(suffix, "down", provider), "litellm_params": _PROVIDERS[provider](f"{wire}/down")} + for provider in ("openai", "anthropic") + ) + + +def _rerouting_models(wire: str, suffix: str) -> tuple[dict[str, JsonValue], ...]: + lists_tools: Final = _owned_name(suffix, "lists_tools", "openai") + return ( + *_down_models(wire, suffix), + { + "model_name": _owned_name(suffix, "auto", "openai"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": lists_tools, + "complexity_router_config": { + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), lists_tools) + }, + }, + }, + { + "model_name": _owned_name(suffix, "direct", "openai"), + "litellm_params": { + **_PROVIDERS["openai"](wire), + "model": f"openai/{_direct(suffix)}", + "additional_drop_params": ["tools"], + }, + }, + { + "model_name": f"{_wildcard(suffix)}/*", + "litellm_params": {**_PROVIDERS["openai"](wire), "model": "openai/*", "additional_drop_params": ["tools"]}, + }, + ) + + +def _rerouting_settings(suffix: str) -> dict[str, JsonValue]: + return { + "fallbacks": [ + {_owned_name(suffix, "down", provider): [_owned_name(suffix, "lists_tools", provider)]} + for provider in ("openai", "anthropic") + ], + "model_group_alias": { + _alias(suffix, variant): _owned_name(suffix, variant, "openai") for variant in ("lists_tools", "plain") + }, + "routing_groups": [ + { + "group_name": _routing_group(suffix, provider), + "models": [_owned_name(suffix, "lists_tools", provider)], + "routing_strategy": "simple-shuffle", + } + for provider in ("openai", "anthropic") + ], + "num_retries": 0, + "disable_cooldowns": True, + } + + +@pytest.fixture(scope="module") +def rerouting(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-rerouting") + setup: Final = _OwnedSetup(router_settings=_rerouting_settings, extra_models=_rerouting_models) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.fixture(scope="module") +def default_fallbacks(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Owned]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-default-fallbacks") + setup: Final = _OwnedSetup( + router_settings=lambda suffix: { + "default_fallbacks": [_owned_name(suffix, "lists_tools", "openai")], + "num_retries": 0, + "disable_cooldowns": True, + }, + extra_models=_down_models, + ) + with _owned_proxy(rig, directory, setup) as owned: + yield owned + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("checked", (_CHECKED["chat"], _CHECKED["messages"]), ids=("chat", "messages")) +def test_a_router_fallback_to_a_group_listing_tools_answers_a_bad_tool_entry( + rerouting: _Owned, checked: _Checked +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, "down", checked.provider) + response: Final = checked.send(rerouting.rig, model, marker, {"tools": [checked.tool, checked.bad_entry]}) + _assert_failed_over(rerouting.rig, checked, response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize( + ("variant", "fields"), + [("plain", {}), ("down", {"fallbacks": []})], + ids=["group_outside_the_chain", "chain_cleared_by_the_request"], +) +def test_a_router_fallback_that_cannot_apply_still_rejects_a_bad_tool_entry( + rerouting: _Owned, variant: str, fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, variant, "openai") + response: Final = _send(rerouting.rig, CHAT, marker, False, model=model, tools=[CHAT.tool, 5], **fields) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_TOKEN_COUNTS +def test_token_counting_on_a_group_with_a_router_fallback_rejects_a_bad_tool_entry( + rerouting: _Owned, checked: _Checked +) -> None: + marker: Final = _marker() + model: Final = _owned_name(rerouting.suffix, "down", "anthropic") + response: Final = checked.send(rerouting.rig, model, marker, {"tools": [checked.tool, 5]}) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_an_alias_of_a_group_listing_tools_answers_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=_alias(rerouting.suffix, "lists_tools"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_an_alias_of_a_plain_group_rejects_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, CHAT, marker, False, model=_alias(rerouting.suffix, "plain"), tools=[CHAT.tool, 5] + ) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_routing_group_over_a_group_listing_tools_answers_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=_routing_group(rerouting.suffix, "openai"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_count_tokens_by_a_routing_group_rejects_a_bad_tool_entry(rerouting: _Owned) -> None: + marker: Final = _marker() + response: Final = _count_tokens( + rerouting.rig, _routing_group(rerouting.suffix, "anthropic"), marker, {"tools": [MESSAGES.tool, 5]} + ) + _assert_held_back(rerouting.rig, response, marker) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize( + ("model", "fields"), + [ + (lambda suffix: _owned_name(suffix, "auto", "openai"), {}), + (lambda suffix: f"{_wildcard(suffix)}/{_OPENAI_BACKEND}", {}), + (lambda suffix: f"openai/{_direct(suffix)}", {}), + (lambda suffix: f"openai/{_direct(suffix)}", {"specific_deployment": True}), + ], + ids=["auto_router", "wildcard_route", "model_string", "specific_deployment"], +) +def test_a_model_the_router_resolves_later_answers_a_bad_tool_entry_its_deployment_list_drops( + rerouting: _Owned, model: Callable[[str], str], fields: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + response: Final = _send( + rerouting.rig, + CHAT, + marker, + False, + model=model(rerouting.suffix), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + **fields, + ) + _assert_forwarded(rerouting.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_default_fallbacks_to_a_group_listing_tools_answer_a_bad_tool_entry(default_fallbacks: _Owned) -> None: + marker: Final = _marker() + response: Final = _send( + default_fallbacks.rig, + CHAT, + marker, + False, + model=_owned_name(default_fallbacks.suffix, "down", "openai"), + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + ) + _assert_failed_over(default_fallbacks.rig, _CHECKED["chat"], response, marker, ()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_a_user_config_model_named_like_a_proxy_group_drops_tools_its_own_list_names(client_side: _Owned) -> None: + marker: Final = _marker() + name: Final = _owned_name(client_side.suffix, "plain", "openai") + user_config: Final[dict[str, JsonValue]] = { + "model_list": [ + { + "model_name": name, + "litellm_params": { + **_PROVIDERS["openai"](client_side.rig.wire.url), + "additional_drop_params": ["tools"], + }, + } + ], + } + response: Final = _send( + client_side.rig, + CHAT, + marker, + False, + model=name, + tools=[CHAT.tool, _CHAT_BAD_ENTRY], + user_config=user_config, + ) + _assert_forwarded(client_side.rig, _CHECKED["chat"], response, marker, ()) + + +_GUARDED_NAMES: Final = ("chat", "messages", "responses", "generate_content") +_EACH_GUARDED: Final = pytest.mark.parametrize( + "checked", tuple(_CHECKED[name] for name in _GUARDED_NAMES), ids=_GUARDED_NAMES +) +_GUARDRAIL_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api" + + +@dataclass(frozen=True, slots=True) +class _Guarded: + rig: _Rig + vendor: Wire + models: Mapping[str, str] + + +def _guarded_config(path: Path, wire: str, vendor: str, models: Mapping[str, str]) -> Path: + path.write_text( + json.dumps( + { + "model_list": [ + {"model_name": name, "litellm_params": _PROVIDERS[provider](wire)} + for provider, name in models.items() + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "guardrails": [ + { + "guardrail_name": "malformed-tools-vendor", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": vendor, + "api_key": "synthetic-guardrail-key", + }, + } + ], + } + ) + ) + return path + + +def _guardrail_verdict(request: Request) -> Reply: + assert request.target == _GUARDRAIL_VENDOR_PATH, request.target + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +@pytest.fixture(scope="module") +def guarded(rig: _Rig, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Guarded]: + directory: Final = tmp_path_factory.mktemp("malformed-tools-guarded") + suffix: Final = uuid.uuid4().hex[:12] + models: Final = MappingProxyType( + {provider: f"malformed-tools-guarded-{provider}-{suffix}" for provider in _PROVIDERS} + ) + with wire_server(_guardrail_verdict) as vendor: + config: Final = _guarded_config(directory / "guarded.yaml", rig.wire.url, vendor.url, models) + with owned_proxy_process(rig.gateway, directory, {}, config=config, workers=2) as proxy: + yield _Guarded(replace(rig, gateway=proxy.gateway), vendor, models) + + +def _vendor_payloads(guarded: _Guarded, marker: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + _JSON_OBJECT.validate_json(request.body) + for request in guarded.vendor.drain() + if marker in request.body.decode() + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@_EACH_GUARDED +def test_a_pre_call_guardrail_passes_well_formed_tools_through_to_the_provider( + guarded: _Guarded, checked: _Checked +) -> None: + marker: Final = _marker() + response: Final = checked.send(guarded.rig, guarded.models[checked.provider], marker, {"tools": [checked.tool]}) + _assert_forwarded(guarded.rig, checked, response, marker, ("get_weather",)) + payloads: Final = _vendor_payloads(guarded, marker) + assert len(payloads) == 1, payloads + assert all(payload.get("tools") is None or isinstance(payload["tools"], list) for payload in payloads), payloads + + +_PROVIDER_WIRES: Final[Mapping[str, tuple[_Family, Callable[[bool], str]]]] = MappingProxyType( + { + "openai": (CHAT, lambda stream: "/v1/chat/completions"), + "openai_responses": (RESPONSES, lambda stream: "/v1/responses"), + "anthropic": (MESSAGES, lambda stream: "/v1/messages"), + "gemini": (GENERATE, lambda stream: ":streamGenerateContent" if stream else ":generateContent"), + } +) +_NATIVE_PROVIDERS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "chat": ("openai",), + "messages": ("anthropic",), + "responses": ("openai", "openai_responses"), + "generate_content": ("gemini",), + } +) +_ROUTED_ELSEWHERE: Final[Mapping[tuple[str, str], str]] = MappingProxyType({("messages", "openai"): "openai_responses"}) +_BRIDGES: Final = tuple( + (family, provider) + for family, provider in product(_FAMILIES, _PROVIDER_WIRES) + if provider not in _NATIVE_PROVIDERS[family.name] +) + + +@pytest.fixture(scope="module") +def bridge_groups(rig: _Rig) -> Iterator[Mapping[str, str]]: + with rig.gateway.scenario() as scenario: + responses_group: Final = scenario.model( + model=f"openai/responses/{_OPENAI_BACKEND}", api_base=f"{rig.wire.url}/v1", api_key=_PROVIDER_KEY + ) + wait_for_worker_readiness(rig.gateway, responses_group, 1, _PROXY_WORKERS) + yield MappingProxyType( + { + "openai": rig.openai, + "openai_responses": responses_group, + "anthropic": rig.anthropic, + "gemini": rig.gemini, + } + ) + + +@_STREAMS +@pytest.mark.parametrize( + ("family", "provider"), _BRIDGES, ids=tuple(f"{family.name}_to_{provider}" for family, provider in _BRIDGES) +) +def test_well_formed_tools_cross_the_bridge_to_another_provider( + rig: _Rig, bridge_groups: Mapping[str, str], family: _Family, provider: str, stream: bool +) -> None: + marker: Final = _marker() + response: Final = _send(rig, family, marker, stream, model=bridge_groups[provider], tools=[family.tool]) + assert response.status_code == 200, response.text + assert _answer(_tag(marker)) in response.text, response.text + named_as, target = _PROVIDER_WIRES[_ROUTED_ELSEWHERE.get((family.name, provider), provider)] + assert _tool_names(named_as, _one_delivery(rig, marker, target(stream))) == ("get_weather",) diff --git a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py index 920002316bc..00f2c430ea5 100644 --- a/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py +++ b/tests/integration/providers/test_openai_dialect_prompt_cache_breakpoint_owned_proxy.py @@ -2,8 +2,8 @@ import asyncio import re import threading from collections.abc import Callable, Iterator -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, replace from pathlib import Path from queue import SimpleQueue from types import MappingProxyType @@ -15,7 +15,7 @@ import psutil import pytest from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment from integration._support.database import read_rows -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.responses_vendor import newest_marker from integration._support.wire import Reply, Request, Wire, wire_server from integration.providers._cache_control_marks_support import owned_config @@ -27,6 +27,7 @@ from integration.providers._mantle_gpt_prompt_cache_support import ( SYSTEM_POINT, TOKEN, Outcome, + Step, assert_answered, assert_burst_landed, assert_wire, @@ -60,7 +61,6 @@ _CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) _RESTART_TIMEOUT: Final = int(4 * graceful_stop_seconds() + 240) _HELD_BURST: Final = 20 _SPREAD_BURST: Final = 12 -_SPREAD_ATTEMPTS: Final = 20 _MANTLE_NAME: Final = "mantle-gpt-cache-owned" @@ -122,7 +122,9 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: release.set() held: Final[SimpleQueue[str]] = SimpleQueue() with gateway_from_environment() as environment, wire_server(_peer(release, held)) as wire: - with owned_proxy_process(environment, directory, _overrides(wire), workers=2) as owned: + with owned_proxy_process( + environment, directory, _overrides(wire), workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: eventually(lambda: len(_started_workers(owned.log)), lambda count: count == 2, seconds=60) wire.drain() yield _Rig(owned.gateway, wire, owned, release, held) @@ -204,38 +206,48 @@ def test_b11_foundry_gpt6_chat_stays_on_the_plain_chat_wire(rig: _Rig) -> None: success_row(name, marker) -def _spread_burst(rig: _Rig, name: str, workers: tuple[int, ...]) -> dict[int, int]: - plan: Final = plan_burst(_SPREAD_BURST) - _drain(rig.held) - rig.wire.drain() - rig.release.clear() - with ThreadPoolExecutor(max_workers=_SPREAD_BURST) as pool: - futures: Final = tuple( - pool.submit(send, rig.gateway, endpoint, request_body(endpoint, name, prompt_text(marker), stream=stream)) - for endpoint, stream, marker in plan - ) - eventually(rig.held.qsize, lambda size: size == _SPREAD_BURST, seconds=60) - held_by: Final = {pid: _open_peer_connections(pid, rig.wire.url) for pid in workers} - rig.release.set() - outcomes: Final = tuple(future.result() for future in futures) - assert sum(held_by.values()) == _SPREAD_BURST, held_by - assert_burst_landed( - rig.wire, name, tuple((marker, outcome) for (_, _, marker), outcome in zip(plan, outcomes)), marked=True - ) - return held_by - - @pytest.mark.timeout(_CELL_TIMEOUT) def test_e6_every_worker_of_a_two_worker_proxy_marks_the_mantle_request(rig: _Rig) -> None: - with rig.gateway.scenario() as scenario: + plan: Final = plan_burst(_SPREAD_BURST) + half: Final = len(plan) // 2 + with ( + rig.gateway.scenario() as scenario, + httpx.Client( + base_url=str(rig.gateway.client.base_url), + timeout=60, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as unpooled, + ThreadPoolExecutor(max_workers=len(plan)) as pool, + ): name: Final = mantle_deployment(rig.gateway, scenario, rig.wire) - workers: Final = _live_workers(rig.owned.log) - assert len(workers) == 2, workers - for _ in range(_SPREAD_ATTEMPTS): - if all(count > 0 for count in _spread_burst(rig, name, workers).values()): - break - else: - raise AssertionError("One worker never took a marked request") + first, second = sorted(_live_workers(rig.owned.log)) + gateway: Final = replace(rig.gateway, client=unpooled) + _drain(rig.held) + rig.wire.drain() + rig.release.clear() + + def submit(part: tuple[Step, ...]) -> tuple[Future[Outcome], ...]: + return tuple( + pool.submit(send, gateway, endpoint, request_body(endpoint, name, prompt_text(marker), stream=stream)) + for endpoint, stream, marker in part + ) + + try: + with rig.owned.suspended(second): + on_first: Final = submit(plan[:half]) + eventually(rig.held.qsize, lambda size: size == half, seconds=60) + with rig.owned.suspended(first): + on_second: Final = submit(plan[half:]) + eventually(rig.held.qsize, lambda size: size == len(plan), seconds=60) + held_by: Final = {pid: _open_peer_connections(pid, rig.wire.url) for pid in (first, second)} + finally: + rig.release.set() + outcomes: Final = tuple(future.result() for future in (*on_first, *on_second)) + assert held_by == {first: half, second: len(plan) - half}, held_by + assert_burst_landed( + rig.wire, name, tuple((marker, outcome) for (_, _, marker), outcome in zip(plan, outcomes)), marked=True + ) def _mantle_config(wire: Wire, directory: Path, **litellm_params: JsonValue) -> Path: diff --git a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py index 0101824054a..6beccefe58f 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py @@ -529,7 +529,7 @@ def _marked_posts(received: tuple[Request, ...], marker: str) -> tuple[Request, @pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) -def test_a_previous_response_id_turn_replays_the_function_call_without_the_search_call(owned_rig: _OwnedRig) -> None: +def test_a_previous_response_id_turn_replays_the_search_call_beside_the_function_call(owned_rig: _OwnedRig) -> None: marker: Final = uuid.uuid4().hex owned_rig.wire.drain() proxy: Final = owned_rig.proxy.gateway @@ -566,11 +566,8 @@ def test_a_previous_response_id_turn_replays_the_function_call_without_the_searc _assert_forwarded_with_tool_search(opening) _assert_forwarded_with_tool_search(replay) messages: Final = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(replay.body)["messages"]) - assert "srvtoolu_" not in json.dumps(messages) and "tool_search" not in json.dumps(messages), messages (assistant,) = [turn for turn in messages if turn["role"] == "assistant"] - blocks: Final = rv.ITEMS.validate_python(assistant["content"]) - tool_uses: Final = [(block["type"], block.get("id")) for block in blocks if block["type"] == "tool_use"] - assert tool_uses == [("tool_use", f"toolu_{marker}")], blocks + assert assistant["content"] == _message(marker)["content"], assistant assert [block["tool_use_id"] for block in _tool_results(messages)] == [f"toolu_{marker}"], messages diff --git a/tests/integration/routing/test_messages_stream_retry_budget_sources_owned_proxy.py b/tests/integration/routing/test_messages_stream_retry_budget_sources_owned_proxy.py index 93e4f97ae1e..0257b8372b0 100644 --- a/tests/integration/routing/test_messages_stream_retry_budget_sources_owned_proxy.py +++ b/tests/integration/routing/test_messages_stream_retry_budget_sources_owned_proxy.py @@ -1,7 +1,9 @@ import json +import threading import uuid from collections.abc import Callable, Iterator from dataclasses import dataclass +from itertools import groupby, islice from pathlib import Path from typing import Final, Literal @@ -12,6 +14,7 @@ from integration._support.anthropic_sse import ( LIFECYCLE, PING, Attempts, + SseEvent, delta_text, dropping_reply, error_body, @@ -55,7 +58,9 @@ def _served_id(backend: str, marker: str, attempt: int) -> str: return f"msg_{backend}_{marker}_a{attempt}" -def _primary_reply(behavior: str, attempt: int, full: tuple[bytes, bytes, bytes, bytes]) -> Reply: +def _primary_reply( + behavior: str, attempt: int, full: tuple[bytes, bytes, bytes, bytes], content_gate: threading.Event +) -> Reply: dropped: Final = dropping_reply(full, abort_after=1) match behavior: case "drop-once": @@ -63,7 +68,13 @@ def _primary_reply(behavior: str, attempt: int, full: tuple[bytes, bytes, bytes, case "drop-always": return dropped case "hold-ping": - return stream_reply((full[0], PING, full[1] + full[2] + full[3]), pause=0.25) + return Reply( + content_type="text/event-stream", + chunks=(full[0], full[1] + full[2] + full[3]), + gate_after_first=content_gate, + heartbeat=PING, + pause_between_chunks=0.25, + ) case "drop-then-too-long": if attempt == 1: return dropped @@ -81,7 +92,7 @@ def _fallback_reply(behavior: str, full: tuple[bytes, bytes, bytes, bytes]) -> R return stream_reply(full) -def _respond(attempts: Attempts) -> Callable[[Request], Reply]: +def _respond(attempts: Attempts, content_gate: threading.Event) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: assert (request.method, request.target) == ("POST", "/v1/messages"), request assert request.headers["x-api-key"] == _API_KEY, request.headers @@ -93,11 +104,21 @@ def _respond(attempts: Attempts) -> Callable[[Request], Reply]: full: Final = message_stream(_served_id(backend, marker, attempt), backend, _TEXT) if backend != _PRIMARY: return _fallback_reply(behavior, full) - return _primary_reply(behavior, attempt, full) + return _primary_reply(behavior, attempt, full, content_gate) return respond +def _stream_body(model: str, behavior: Behavior, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": f"{behavior}:{marker}"}], + **extra, + } + + def _deployment(name: str, backend: str, **extra: JsonValue) -> dict[str, JsonValue]: return { "model_name": name, @@ -137,19 +158,10 @@ def _config(wire: Wire, directory: Path) -> Path: class _Rig: proxy: Gateway attempts: Attempts + content_gate: threading.Event def stream(self, model: str, behavior: Behavior, marker: str, **extra: JsonValue) -> httpx.Response: - return self.proxy.request( - "POST", - "/v1/messages", - { - "model": model, - "max_tokens": 16, - "stream": True, - "messages": [{"role": "user", "content": f"{behavior}:{marker}"}], - **extra, - }, - ) + return self.proxy.request("POST", "/v1/messages", _stream_body(model, behavior, marker, **extra)) def attempts_on(self, backend: str, marker: str) -> int: return self.attempts.count(f"{backend}:{marker}") @@ -158,10 +170,14 @@ class _Rig: @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: attempts: Final = Attempts() + content_gate: Final = threading.Event() directory: Final = tmp_path_factory.mktemp("messages-retry-budget-sources") - with gateway_from_environment() as gateway, wire_server(_respond(attempts)) as wire: + with gateway_from_environment() as gateway, wire_server(_respond(attempts, content_gate)) as wire: with owned_proxy(gateway, directory, {}, config=_config(wire, directory), workers=2) as proxy: - yield _Rig(proxy, attempts) + try: + yield _Rig(proxy, attempts, content_gate) + finally: + content_gate.set() def _assert_completed(response: httpx.Response, served_id: str) -> None: @@ -202,12 +218,32 @@ def test_request_num_retries_zero_turns_the_router_budget_off(rig: _Rig) -> None assert rig.attempts_on(_PRIMARY, marker) == 1 +def _live_events(lines: Iterator[str]) -> Iterator[SseEvent]: + for has_text, block in groupby(lines, key=bool): + if has_text: + yield from parse_sse("\n".join(block)) + + def test_lifecycle_frames_are_held_until_content_while_pings_go_out_live(rig: _Rig) -> None: marker: Final = uuid.uuid4().hex - response: Final = rig.stream(_ROUTER_BUDGET, "hold-ping", marker) - assert response.status_code == 200, response.text - events: Final = parse_sse(response.text) - assert event_types(events) == ("ping", *LIFECYCLE), events + rig.content_gate.clear() + try: + with rig.proxy.client.stream( + "POST", + "/v1/messages", + json=_stream_body(_ROUTER_BUDGET, "hold-ping", marker), + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) as response: + assert response.status_code == 200, response.read() + arriving: Final = _live_events(response.iter_lines()) + before_content: Final = tuple(islice(arriving, 2)) + rig.content_gate.set() + after_content: Final = tuple(arriving) + finally: + rig.content_gate.set() + assert event_types(before_content) == ("ping", "ping"), before_content + events: Final = before_content + after_content + assert tuple(kind for kind in event_types(events) if kind != "ping") == LIFECYCLE, events assert message_id(events) == _served_id(_PRIMARY, marker, 1), events assert delta_text(events) == _TEXT, events assert rig.attempts_on(_PRIMARY, marker) == 1 diff --git a/tests/integration/routing/test_user_config_routing.py b/tests/integration/routing/test_user_config_routing.py index 1c50a78201b..af3df8a38ff 100644 --- a/tests/integration/routing/test_user_config_routing.py +++ b/tests/integration/routing/test_user_config_routing.py @@ -3,8 +3,10 @@ import uuid from pathlib import Path from typing import Final +import pytest + from integration._support.client import Gateway -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server USER_KEY: Final = "sk-user-supplied-" + uuid.uuid4().hex @@ -30,18 +32,13 @@ def _completion(request: Request) -> Reply: ) +def _deployment_params(upstream_url: str) -> dict[str, object]: + return {"model": "openai/gpt-4.1-mini", "api_base": upstream_url + "/v1", "api_key": USER_KEY} + + 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, - }, - } - ], + "model_list": [{"model_name": "user-config-deployment", "litellm_params": _deployment_params(upstream_url)}], "num_retries": 0, } @@ -95,6 +92,54 @@ def test_user_config_routes_to_the_user_supplied_deployment_when_opted_in(gatewa assert json.loads(completions[0].body)["model"] == "gpt-4.1-mini" +def _identity_in_deployment(upstream_url: str, identity_param: str) -> dict[str, object]: + return { + "model_list": [ + { + "model_name": "user-config-deployment", + "litellm_params": {**_deployment_params(upstream_url), identity_param: "caller-chosen"}, + } + ], + "num_retries": 0, + } + + +def _identity_in_defaults(upstream_url: str, identity_param: str) -> dict[str, object]: + return {**_user_config(upstream_url), "default_litellm_params": {identity_param: "caller-chosen"}} + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_user_config_cannot_carry_a_server_owned_identity_param_even_when_opted_in( + gateway: Gateway, tmp_path: Path +) -> None: + cases: Final = ( + (_identity_in_deployment, "openai_identity_provider_id"), + (_identity_in_defaults, "openai_service_account_id"), + ) + 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: + responses: Final = tuple( + ( + identity_param, + candidate.request( + "POST", + "/v1/chat/completions", + { + "model": "user-config-deployment", + "messages": [{"role": "user", "content": "user config identity"}], + "user_config": placement(upstream.url, identity_param), + }, + ), + ) + for placement, identity_param in cases + ) + for identity_param, response in responses: + assert response.status_code == 401, response.text + assert identity_param in response.text, response.text + assert tuple(request for request in upstream.drain() if request.target == "/v1/chat/completions") == () + + 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)) diff --git a/tests/integration/routing/test_wildcard_model_listing.py b/tests/integration/routing/test_wildcard_model_listing.py new file mode 100644 index 00000000000..967a515ec08 --- /dev/null +++ b/tests/integration/routing/test_wildcard_model_listing.py @@ -0,0 +1,71 @@ +import uuid +from collections import Counter +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(180) + +_PRICED_DEPLOYMENT_ID: Final = f"priced-deployment-{uuid.uuid4().hex}" +_REGISTERED_AI21_CHAT_MODEL: Final = f"jamba-house-{uuid.uuid4().hex}" + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("wildcard-listing") + with gateway_from_environment() as base: + credentials: Final = {"api_base": f"{base.upstream_url}/v1", "api_key": "synthetic-listing-key"} + config: Final = { + **yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()), + "model_list": [ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", **credentials}}, + { + "model_name": "priced-chat", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + **credentials, + }, + "model_info": {"id": _PRICED_DEPLOYMENT_ID}, + }, + {"model_name": "ai21/*", "litellm_params": {"model": "ai21/*", **credentials}}, + { + "model_name": "house-jamba", + "litellm_params": {"model": f"ai21/{_REGISTERED_AI21_CHAT_MODEL}", **credentials}, + "model_info": {"litellm_provider": "ai21", "mode": "chat"}, + }, + ], + } + path: Final = directory / "wildcard-listing.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(base, directory, {}, config=path) as proxy: + yield proxy + + +def _listed_ids(gateway: Gateway) -> list[JsonValue]: + data: Final = gateway.get("/v1/models")["data"] + assert isinstance(data, list), data + return [entry["id"] for entry in data if isinstance(entry, dict)] + + +def test_wildcard_listing_after_a_priced_call_names_each_model_once_and_no_deployment_id(candidate: Gateway) -> None: + candidate.chat("priced-chat") + + listed: Final = _listed_ids(candidate) + + assert "openai/gpt-4o-mini" in listed and "priced-chat" in listed, listed + assert {model: count for model, count in Counter(listed).items() if count > 1} == {}, listed + assert [model for model in listed if _PRICED_DEPLOYMENT_ID in str(model)] == [], listed + + +def test_ai21_wildcard_listing_names_a_chat_model_registered_from_the_config(candidate: Gateway) -> None: + listed: Final = _listed_ids(candidate) + + assert f"ai21/{_REGISTERED_AI21_CHAT_MODEL}" in listed, listed diff --git a/tests/integration/spend/test_disable_end_user_cost_tracking.py b/tests/integration/spend/test_disable_end_user_cost_tracking.py index f3ec572b69a..d0233378d1e 100644 --- a/tests/integration/spend/test_disable_end_user_cost_tracking.py +++ b/tests/integration/spend/test_disable_end_user_cost_tracking.py @@ -695,6 +695,7 @@ def test_flag_on_burst_lands_each_row_once_through_sink_outage_and_worker_kill(r before_kill: Final = _burst(rig, models, key, calls[: len(_SURFACES) * 2]) _assert_burst_landed_once(rig, models, before_kill) assert not rig.lens.outages.acquire(blocking=False), "the Lens sink outage was never hit" + _team_requests_flushed(team, len(before_kill)) victim: Final = min(rig.proxy.worker_pids()) replacement: Final = rig.proxy.kill_worker(victim) assert replacement != victim, (victim, replacement) diff --git a/tests/integration/streaming/test_faked_responses_stream_chaos.py b/tests/integration/streaming/test_faked_responses_stream_chaos.py index 28efe912157..fca99c9d525 100644 --- a/tests/integration/streaming/test_faked_responses_stream_chaos.py +++ b/tests/integration/streaming/test_faked_responses_stream_chaos.py @@ -24,7 +24,7 @@ from integration._support.client import ( string_value, ) from integration._support.database import read_rows -from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from integration._support.ports import allocate_port from integration.streaming.faked_responses_case import ( @@ -111,7 +111,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: port: Final = allocate_port() with ( gateway_from_environment() as shared, - owned_proxy_process(shared, directory, {}, config=_config(port, directory), workers=2) as owned, + owned_proxy_process( + shared, + directory, + {}, + config=_config(port, directory), + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + ) as owned, ): yield _Rig(port, owned) @@ -306,9 +313,11 @@ class _Held: return _respond(request) -async def _held_burst(gateway: Gateway, calls: tuple[_Call, ...], held: _Held) -> asyncio.Task[tuple[_Served, ...]]: +async def _held_burst( + gateway: Gateway, calls: tuple[_Call, ...], held: _Held, *, already_held: int = 0 +) -> asyncio.Task[tuple[_Served, ...]]: burst: Final = asyncio.create_task(_burst(gateway, calls, tolerate_transport_errors=True)) - await asyncio.to_thread(eventually, held.markers.qsize, lambda size: size == len(calls), 60) + await asyncio.to_thread(eventually, held.markers.qsize, lambda size: size == already_held + len(calls), 60) return burst @@ -319,12 +328,18 @@ def _held_by(rig: _Rig, upstream: str) -> Mapping[int, int]: async def test_held_mixed_burst_spreads_over_both_workers_and_each_value_arrives_once(rig: _Rig) -> None: calls: Final = _calls(30) held: Final = _Held(threading.Event(), SimpleQueue()) + first, second = sorted(rig.proxy.worker_pids()) with wire_server(held.respond, port=rig.port) as wire: - burst: Final = await _held_burst(rig.gateway, calls, held) - held_by: Final = _held_by(rig, wire.url) - held.release.set() - served: Final = await burst - assert len(held_by) == 2 and all(held_by.values()) and sum(held_by.values()) == 30, held_by + try: + with rig.proxy.suspended(second): + on_first: Final = await _held_burst(rig.gateway, calls[:15], held) + with rig.proxy.suspended(first): + on_second: Final = await _held_burst(rig.gateway, calls[15:], held, already_held=15) + held_by: Final = _held_by(rig, wire.url) + finally: + held.release.set() + served: Final = await on_first + await on_second + assert held_by == {first: 15, second: 15}, held_by assert len(served) == 30 for item in served: _assert_each_value_once(item) diff --git a/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py b/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py index 85017667c74..796e706c569 100644 --- a/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py +++ b/tests/integration/streaming/test_google_genai_stream_usage_settings_wire.py @@ -1,21 +1,19 @@ import json -import re import threading import uuid from collections.abc import Callable, Mapping, Sequence -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import ExitStack from hashlib import sha256 from pathlib import Path from typing import Final import httpx -import psutil import pytest import yaml from integration._support.client import Gateway, eventually, held, object_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -38,7 +36,6 @@ _GEMINI_USAGE: Final = { "totalTokenCount": _PROMPT_TOKENS + _COMPLETION_TOKENS, } _ZERO_USAGE: Final = {"promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0} -_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _BURST_SELECT: Final = ( 'SELECT api_key, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs"' " WHERE api_key = ANY(string_to_array(%s, ','))" @@ -67,6 +64,7 @@ _UPSTREAM_ASKS_FOR_USAGE: Final = { }, } _SETTING_VALUES: Final = {"always_off": False, "always_on": True} +_STREAMS_PER_WORKER: Final = 6 def _openai_frame(chunk: Mapping[str, JsonValue], response_id: str) -> bytes: @@ -273,20 +271,15 @@ def _stream_or_dropped(client: httpx.Client, key: str) -> httpx.Response | None: return None -def _local_port(client: httpx.Client) -> int: - with client.stream("GET", "/health/liveliness") as response: - port: Final = int(response.extensions["network_stream"].get_extra_info("client_addr")[1]) - response.read() - assert response.status_code == 200, response.text - return port - - -def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: - return frozenset( - connection.raddr.port - for connection in psutil.Process(pid).net_connections(kind="tcp") - if connection.raddr and connection.laddr.port == proxy_port +def _held_streams( + stack: ExitStack, pool: ThreadPoolExecutor, url: str, keys: Sequence[str], wire: Wire, held: int +) -> tuple[Future[httpx.Response | None], ...]: + clients: Final = tuple(stack.enter_context(httpx.Client(base_url=url, timeout=60, trust_env=False)) for _ in keys) + futures: Final = tuple( + pool.submit(_stream_or_dropped, client, key) for client, key in zip(clients, keys, strict=True) ) + eventually(wire.received.qsize, lambda size: size >= held, seconds=60) + return futures @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -296,35 +289,27 @@ def test_worker_killed_mid_stream_leaves_the_sibling_billing_each_stream_once(ga wire: Final = stack.enter_context(wire_server(_peer(gate))) stack.callback(gate.set) config: Final = _owned_config(tmp_path / "worker-kill.yaml", gateway, wire, {}) - owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) - workers: Final = eventually( - lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2, - seconds=30, + owned: Final = stack.enter_context( + owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT) ) + victim, survivor = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=30)) scenario: Final = stack.enter_context(owned.gateway.scenario()) - proxy_url: Final = owned.gateway.client.base_url - clients: Final = tuple( - stack.enter_context(httpx.Client(base_url=str(proxy_url), timeout=60, trust_env=False)) for _ in range(12) - ) - keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in clients) - ports: Final = tuple(_local_port(client) for client in clients) - pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) - futures: Final = tuple( - pool.submit(_stream_or_dropped, client, key) for client, key in zip(clients, keys, strict=True) - ) - eventually(wire.received.qsize, lambda size: size >= len(clients), seconds=60) - shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} - assert sum(map(len, shares.values())) == len(clients), shares - victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + victim_keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in range(_STREAMS_PER_WORKER)) + survivor_keys: Final = tuple(scenario.key(models=[_ADAPTER_MODEL]) for _ in range(_STREAMS_PER_WORKER)) + proxy_url: Final = str(owned.gateway.client.base_url) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=2 * _STREAMS_PER_WORKER)) + with owned.suspended(survivor): + on_victim: Final = _held_streams(stack, pool, proxy_url, victim_keys, wire, _STREAMS_PER_WORKER) + with owned.suspended(victim): + on_survivor: Final = _held_streams(stack, pool, proxy_url, survivor_keys, wire, 2 * _STREAMS_PER_WORKER) with owned.replacing_worker(victim) as kill: kill() gate.set() - results: Final = tuple(future.result(timeout=60) for future in futures) - survivors: Final = tuple( - (key, result) for port, key, result in zip(ports, keys, results, strict=True) if port not in shares[victim] - ) - assert survivors, shares + dropped: Final = tuple(future.result(timeout=60) for future in on_victim) + survivors: Final = tuple( + zip(survivor_keys, (future.result(timeout=60) for future in on_survivor), strict=True) + ) + assert dropped == (None,) * _STREAMS_PER_WORKER, dropped for _, result in survivors: assert result is not None _assert_one_stop_with(result, _GEMINI_USAGE) diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 293c38e9f20..912d4cfe84c 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -1983,6 +1983,20 @@ def test_total_cost_bills_cached_tokens_per_line_at_the_batch_cached_rate(): assert result.cost == pytest.approx(long_line + short_line) +def test_extract_credentials_keeps_openai_workload_identity_params(): + params = { + "openai_identity_provider_id": "idp_deployment", + "openai_service_account_id": "user-deployment", + "openai_identity_token_file": "/var/run/secrets/token", + "model": "gpt-4o", + } + assert bu._extract_file_access_credentials(params) == { + "openai_identity_provider_id": "idp_deployment", + "openai_service_account_id": "user-deployment", + "openai_identity_token_file": "/var/run/secrets/token", + } + + # --------------------------------------------------------------------------- # # batch_cost_is_final # --------------------------------------------------------------------------- # diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 26306a2d39a..2106436edb8 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -379,48 +379,47 @@ def test_a_replacing_block_that_never_kills_its_worker_fails_and_leaves_it_servi assert owned.worker_pids() == frozenset((worker,)) -def _kill_of(owned: OwnedProxy, worker: int) -> WorkerKill: - return WorkerKill(psutil.Process(worker), psutil.Process(owned.process.pid), 5, threading.Event()) +@contextmanager +def _unreaped_child() -> Iterator[psutil.Process]: + with subprocess.Popen((sys.executable, "-I", "-c", "import time; time.sleep(60)")) as child: + try: + yield psutil.Process(child.pid) + finally: + child.kill() -def test_a_worker_kill_returns_once_the_supervisor_reaped_the_victim_and_cannot_fire_twice(tmp_path: Path) -> None: - with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): - kill: Final = _kill_of(owned, worker) +def test_a_worker_kill_returns_once_the_victim_is_dead_before_anything_reaps_it_and_cannot_fire_twice() -> None: + with _unreaped_child() as victim: + kill: Final = WorkerKill(victim, threading.Event()) kill() - assert not psutil.pid_exists(worker) + assert victim.status() == psutil.STATUS_ZOMBIE assert kill.fired.is_set() with pytest.raises(AssertionError, match="already killed"): kill() -def test_a_frozen_victim_stays_stopped_until_the_kill(tmp_path: Path) -> None: - with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): - kill: Final = _kill_of(owned, worker) +def test_a_frozen_victim_stays_stopped_until_the_kill() -> None: + with _unreaped_child() as victim: + kill: Final = WorkerKill(victim, threading.Event()) assert kill.freeze().status() == psutil.STATUS_STOPPED kill() - assert not psutil.pid_exists(worker) + assert victim.status() == psutil.STATUS_ZOMBIE -def test_a_kill_under_a_suspended_supervisor_returns_with_the_victim_left_unreaped(tmp_path: Path) -> None: - with _supervised(tmp_path, WorkerHealthcheck.LENIENT) as (owned, worker): - kill: Final = _kill_of(owned, worker) - with owned.suspended(owned.process.pid): - kill() - assert psutil.Process(worker).status() == psutil.STATUS_ZOMBIE - eventually(lambda: psutil.pid_exists(worker), lambda exists: not exists, seconds=5) +def _kill_after_freezing_the_supervisor(owned: OwnedProxy, worker: int, supervisor: psutil.Process) -> None: + with owned.replacing_worker(worker) as kill: + supervisor.suspend() + kill() -def test_a_kill_fails_when_the_supervisor_never_reaps_the_victim(tmp_path: Path) -> None: +def test_a_replacing_block_fails_when_the_supervisor_never_reaps_the_killed_worker(tmp_path: Path) -> None: with _supervised(tmp_path, WorkerHealthcheck.STRICT) as (owned, worker): unreaping: Final = psutil.Process(owned.process.pid) - kill: Final = WorkerKill(psutil.Process(worker), psutil.Process(os.getpid()), 0.5, threading.Event()) - unreaping.suspend() try: with pytest.raises(AssertionError, match="did not reap"): - kill() + _kill_after_freezing_the_supervisor(owned, worker, unreaping) finally: unreaping.resume() - assert not kill.fired.is_set() def test_a_strict_healthcheck_proxy_refuses_suspension_before_stopping_anything(tmp_path: Path) -> None: diff --git a/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py index fedfc2b0848..8eefee1dfcf 100644 --- a/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py +++ b/tests/unit/integrations/test_prometheus_deployment_state_proxy_rejects.py @@ -9,12 +9,14 @@ response. A failure raised after the router picked a deployment (a post-call guardrail block, a provider error) carries no flag and keeps its deployment labels. """ +from typing import Final + import pytest from fastapi import HTTPException from prometheus_client import REGISTRY from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY -from litellm.exceptions import GuardrailRaisedException +from litellm.exceptions import GuardrailRaisedException, MalformedRequestFieldError from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -99,3 +101,15 @@ def test_failure_after_routing_still_marks_deployment_partial_outage(failure: Ex assert _model_id_values(logger.litellm_deployment_state) == {"dep-1"} assert _model_id_values(logger.litellm_deployment_failure_responses) == {"dep-1"} + + +def test_malformed_field_rejection_after_routing_leaves_deployment_healthy() -> None: + logger: Final = PrometheusLogger() + rejection: Final = MalformedRequestFieldError( + message="tools: Input should be a valid list", model="internal-model", llm_provider="openai" + ) + + logger.set_llm_deployment_failure_metrics(_attributed_failure_kwargs(rejection, rejected_before_routing=False)) + + assert logger.litellm_deployment_state._metrics == {} + assert _model_id_values(logger.litellm_deployment_failure_responses) == {""} diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index c8dbbd00d66..7291c4858f6 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,6 @@ +from typing import Final + import pytest import litellm @@ -49,11 +51,12 @@ def test_web_search_cost_high(): ) -def test_get_web_search_options_preserves_explicit_none_error(): - with pytest.raises(TypeError): - StandardBuiltInToolCostTracking.get_web_search_options( - {"web_search_options": None, "tools": [{"type": "web_search_preview"}]} - ) +def test_get_web_search_options_reads_tools_when_options_are_null() -> None: + options: Final = StandardBuiltInToolCostTracking.get_web_search_options( + {"web_search_options": None, "tools": [{"type": "web_search_preview"}]} + ) + + assert options == {"type": "web_search_preview"} def test_get_web_search_options_accepts_non_list_tool_iterables(): @@ -1070,3 +1073,13 @@ def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252) assert get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) is None + + +@pytest.mark.parametrize( + "kwargs", + [{"tools": 5}, {"web_search_options": None}, {"web_search_options": 5}], + ids=["int_tools", "null_web_search_options", "int_web_search_options"], +) +def test_built_in_tool_lookups_ignore_malformed_request_values(kwargs: dict[str, object]) -> None: + assert StandardBuiltInToolCostTracking.get_web_search_options(kwargs) is None + assert StandardBuiltInToolCostTracking.get_file_search_tool_call(kwargs) is None diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index cf713d3684e..0256ecf24ed 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -10,7 +10,9 @@ import sys import threading from typing import Final +import boto3 import pytest +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, @@ -519,6 +521,47 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco from unittest.mock import MagicMock, patch +class _FakeS3Body: + def __init__(self, content: bytes) -> None: + self.content: Final = content + + def read(self) -> bytes: + return self.content + + +class _FakeS3GetObjectResponse(TypedDict): + Body: ReadOnly[_FakeS3Body] + ETag: NotRequired[ReadOnly[str]] + + +class _FakeS3Client: + def __init__(self, content: bytes) -> None: + self.content: Final = content + self.calls: list[tuple[str, str]] = [] + + def get_object(self, *, Bucket: str, Key: str) -> _FakeS3GetObjectResponse: + self.calls.append((Bucket, Key)) + return {"Body": _FakeS3Body(self.content), "ETag": '"fake-etag"'} + + +def test_get_model_cost_map_fetches_s3_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils import get_model_cost_map as module + + source_info = module._cost_map_source_info + for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): + monkeypatch.setattr(source_info, name, getattr(source_info, name)) + + s3_content: Final = _make_models(GetModelCostMap.get_backup_model_count()) + s3_client: Final = _FakeS3Client(json.dumps(s3_content).encode()) + monkeypatch.setattr(boto3, "client", lambda *args, **kwargs: s3_client) + + cost_map: Final = get_model_cost_map("s3://my-bucket/prices/model_prices.json") + + assert cost_map == s3_content + assert s3_client.calls == [("my-bucket", "prices/model_prices.json")] + assert get_model_cost_map_source_info()["source"] == "remote" + + class _SyncSleepRecorder: """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" @@ -955,6 +998,14 @@ class TestGetModelCostMapFallback: assert result.keys() == _backup_keys() assert source["source"] == "local" + def test_should_fallback_when_fetched_json_is_not_object(self, monkeypatch): + result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.Response(200, content=b"[]")) + + assert fetches == 1 + assert result.keys() == _backup_keys() + assert source["source"] == "local" + assert source["fallback_reason"] == f"Remote fetch failed: expected a JSON object from {_URL}, got list" + def test_should_fallback_when_fetched_map_shrinks_dramatically(self, monkeypatch): """When upstream returns far fewer models than backup, should fall back.""" tiny_map = {f"model-{i}": {"litellm_provider": "test"} for i in range(11)} diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f2d4f0554cf..dc61d9bc63e 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -6187,6 +6187,42 @@ def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_resp assert result.usage.prompt_tokens == 4 # type: ignore[attr-defined] +@pytest.mark.parametrize( + "output", + [[], [{"type": "reasoning", "id": "rs_1", "status": "completed", "summary": []}]], + ids=["empty", "reasoning-only"], +) +def test_degraded_responses_payload_keeps_the_served_service_tier_price( + output: list[dict[str, object]], monkeypatch: pytest.MonkeyPatch +): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + rates: Final = { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "input_cost_per_token_flex": 5e-07, + "output_cost_per_token_flex": 2e-06, + } + monkeypatch.setitem( + litellm.model_cost, "flex-fallback-model", {"litellm_provider": "openai", "mode": "responses", **rates} + ) + degraded: Final = _anthropic_messages_logging_obj()._handle_anthropic_messages_response_logging( + result=ResponsesAPIResponse( + id="resp-flex", + created_at=1700000000, + status="completed", + output=output, + service_tier="flex", + usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150), + ) + ) + + assert degraded.choices == [] or degraded.choices[0].message.content in (None, "") # type: ignore[union-attr] + assert litellm.completion_cost(completion_response=degraded, model="flex-fallback-model") == pytest.approx( + 100 * rates["input_cost_per_token_flex"] + 50 * rates["output_cost_per_token_flex"] + ) + + class _SuccessCapturingLogger(CustomLogger): """Records the success payload. success_payload is populated only in async_log_success_event, so it stays None when the buggy no-op diff --git a/tests/unit/litellm_core_utils/test_tool_fields.py b/tests/unit/litellm_core_utils/test_tool_fields.py new file mode 100644 index 00000000000..23d7de98fbc --- /dev/null +++ b/tests/unit/litellm_core_utils/test_tool_fields.py @@ -0,0 +1,187 @@ +import datetime +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm import Router +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.utils import async_pre_call_deployment_hook + +_GOOD_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} +_MALFORMED_TOOLS: Final = ( + 5, + "get_weather", + ["get_weather"], + {"type": "function"}, + [{"type": "function", "function": {"name": "get_weather", "parameters": "object"}}], +) + + +def _deployment(name: str, **params: object) -> Mapping[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "mock_response": f"served by {name}", **params}, + } + + +def _proxy_request(model: str, route: str = "acompletion", **fields: object) -> Mapping[str, object]: + return { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "proxy_server_request": {"url": "/v1/chat/completions", "body": {}}, + "litellm_logging_obj": Logging( + model=model, + messages=[], + stream=False, + call_type=route, + start_time=datetime.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ), + "num_retries": 0, + **fields, + } + + +@pytest.mark.parametrize("tools", _MALFORMED_TOOLS) +@pytest.mark.asyncio +async def test_malformed_tools_fall_back_to_a_deployment_that_drops_them(tools: object) -> None: + router: Final = Router( + model_list=[_deployment("primary"), _deployment("backup", additional_drop_params=["tools"])], + fallbacks=[{"primary": ["backup"]}], + ) + + response: Final = await router.acompletion(**_proxy_request("primary", tools=tools)) + + assert response.choices[0].message.content == "served by backup" + + +@pytest.mark.asyncio +async def test_fallback_on_unsupported_params_false_keeps_the_400() -> None: + router: Final = Router( + model_list=[_deployment("primary"), _deployment("backup", additional_drop_params=["tools"])], + fallbacks=[{"primary": ["backup"]}], + fallback_on_unsupported_params=False, + ) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("primary", tools=5)) + + +@pytest.mark.asyncio +async def test_a_rejected_request_cools_down_no_deployment() -> None: + cools_on_any_400: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}, "cooldown_time": 120} + router: Final = Router( + model_list=[ + {**_deployment("primary"), "model_info": {"id": "primary", **cools_on_any_400}}, + {**_deployment("backup"), "model_info": {"id": "backup", **cools_on_any_400}}, + ], + fallbacks=[{"primary": ["backup"]}], + ) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("primary", tools=5)) + + primary: Final = await router.acompletion(**_proxy_request("primary")) + backup: Final = await router.acompletion(**_proxy_request("backup")) + assert (primary.choices[0].message.content, backup.choices[0].message.content) == ( + "served by primary", + "served by backup", + ) + + +@pytest.mark.parametrize("tools", _MALFORMED_TOOLS) +@pytest.mark.parametrize( + "deployment", + [_deployment("keeps"), _deployment("keeps", tools=[_GOOD_TOOL])], + ids=["plain", "with_configured_tools"], +) +@pytest.mark.asyncio +async def test_malformed_tools_answer_400_when_the_served_deployment_keeps_tools( + tools: object, deployment: Mapping[str, object] +) -> None: + router: Final = Router(model_list=[deployment]) + + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(**_proxy_request("keeps", tools=tools)) + + assert (raised.value.status_code, raised.value.param, raised.value.type) == ( + 400, + "tools", + "invalid_request_error", + ) + + +@pytest.mark.asyncio +async def test_a_request_drop_params_false_overrides_a_deployment_that_drops() -> None: + router: Final = Router(model_list=[_deployment("drops", drop_params=True)]) + + with pytest.raises(litellm.BadRequestError): + await router.acompletion(**_proxy_request("drops", tools=5, drop_params=False)) + + +@pytest.mark.parametrize( + "route", ["acompletion", "aresponses", "anthropic_messages", "agenerate_content", "acreate_interaction"] +) +@pytest.mark.asyncio +async def test_drop_params_drops_every_malformed_tool_field(route: str) -> None: + request: Final = _proxy_request( + "drops", route, drop_params=True, tools=5, tool_choice="auto", web_search_options="high" + ) + + sent: Final = await async_pre_call_deployment_hook(dict(request), route) + + assert {"tools", "tool_choice", "web_search_options"}.isdisjoint(sent) + assert sent["messages"] == request["messages"] + + +@pytest.mark.asyncio +async def test_additional_drop_params_drops_only_the_malformed_tool_entries() -> None: + request: Final = _proxy_request("drops", additional_drop_params=["tools"], tools=[_GOOD_TOOL, "get_weather"]) + + sent: Final = await async_pre_call_deployment_hook(dict(request), "acompletion") + + assert sent["tools"] == [_GOOD_TOOL] + + +@pytest.mark.asyncio +async def test_an_sdk_call_keeps_the_tools_it_passed() -> None: + tools: Final = (_GOOD_TOOL,) + request: Final = { + key: value for key, value in _proxy_request("sdk", tools=tools).items() if key != "proxy_server_request" + } + + sent: Final = await async_pre_call_deployment_hook(dict(request), "acompletion") + + assert sent["tools"] is tools + + +@pytest.mark.parametrize("schema", ["", [], False, 0, {}, '{"type": "object"}']) +@pytest.mark.parametrize( + "tool", + [ + lambda schema: {"type": "function", "function": {"name": "get_weather", "parameters": schema}}, + lambda schema: {"type": "function", "name": "get_weather", "parameters": schema}, + lambda schema: {"name": "get_weather", "input_schema": schema}, + ], + ids=["chat", "responses", "messages"], +) +@pytest.mark.asyncio +async def test_empty_and_json_string_schemas_reach_the_deployment( + tool: Callable[[object], Mapping[str, object]], schema: object +) -> None: + tools: Final = [tool(schema)] + + sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") + + assert sent["tools"] is tools + + +@pytest.mark.parametrize("tools", [(_GOOD_TOOL,), [{"type": "web_search", "function": None}]], ids=["tuple", "null"]) +@pytest.mark.asyncio +async def test_tools_core_accepts_reach_the_deployment(tools: object) -> None: + sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") + + assert sent["tools"] is tools diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 61efc65f8c2..453429df084 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -553,3 +553,83 @@ def test_client_cache_control_changes_nothing_for_a_deployment_that_does_not_rea ) assert marked == unmarked + + +REASONING_MODEL: Final = "gpt-bridge-reasoning-cost" +REASONING_RATES: Final = { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "output_cost_per_reasoning_token": 3e-05, +} +REASONING_RESPONSE: Final = { + "id": "resp_bridge_reasoning", + "object": "response", + "created_at": 1, + "status": "completed", + "model": REASONING_MODEL, + "output": [ + {"type": "reasoning", "id": "rs_1", "status": "completed", "summary": []}, + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "answer", "annotations": []}], + }, + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": { + "input_tokens": 1240, + "output_tokens": 4040, + "total_tokens": 5280, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 3480}, + }, +} + + +@pytest.mark.asyncio +async def test_recomputed_cost_of_a_served_bridge_response_prices_the_reasoning_tokens( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + monkeypatch.setenv("OPENAI_API_KEY", "sk-bridge-reasoning-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + monkeypatch.setitem( + litellm.model_cost, REASONING_MODEL, {"litellm_provider": "openai", "mode": "responses", **REASONING_RATES} + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").respond(status_code=200, json=REASONING_RESPONSE) + logging_obj: Final = Logging( + model=REASONING_MODEL, + messages=MESSAGES, + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="bridge-reasoning-0000-4000-8000-000000000001", + function_id="1234", + ) + + served: Final = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=412, + messages=MESSAGES, + model=f"openai/{REASONING_MODEL}", + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + ) + assert isinstance(served, dict) and "reasoning" not in json.dumps(served["usage"]) + + cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(served), logging_obj=logging_obj + ) + + expected_reasoning: Final = 3480 * REASONING_RATES["output_cost_per_reasoning_token"] + expected_output: Final = expected_reasoning + (4040 - 3480) * REASONING_RATES["output_cost_per_token"] + assert cost == pytest.approx(1240 * REASONING_RATES["input_cost_per_token"] + expected_output) + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output) + assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index 9875cf439cf..f753fb6b8f9 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -753,3 +753,33 @@ def test_azure_deployment_id_selects_chat_completions_endpoint(respx_mock: MockR assert route.call_count == 1 assert route.calls[0].request.url.path == "/openai/deployments/deployment-gpt-4o/chat/completions" assert response.choices[0].message.content == "Hello" + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_leaves_the_chat_route_and_body_untouched(respx_mock: MockRouter) -> None: + url: Final = "https://example-resource.openai.azure.com/openai/deployments/gpt-5.4-nano/chat/completions" + payload: Final = { + "id": "chatcmpl-dated", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + } + route: Final = respx_mock.post(url__startswith=url).mock(return_value=httpx.Response(200, json=payload)) + messages: Final = [{"role": "user", "content": "Hello"}] + + response: Final = litellm.completion( + model="azure/gpt-5.4-nano", + messages=messages, + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-12-01-preview", + responses_api_version="preview", + ) + request: Final = route.calls[0].request + + assert route.call_count == 1 + assert request.url.params["api-version"] == "2024-12-01-preview" + assert json.loads(request.content) == {"model": "gpt-5.4-nano", "messages": messages} + assert response.choices[0].message.content == "Hello" diff --git a/tests/unit/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py index fea5304a689..9b4c618c78b 100644 --- a/tests/unit/llms/azure/response/test_azure_transformation.py +++ b/tests/unit/llms/azure/response/test_azure_transformation.py @@ -1,11 +1,15 @@ import json +from collections.abc import Callable from copy import deepcopy from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +import respx import litellm +from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.azure.responses.o_series_transformation import ( @@ -282,16 +286,25 @@ class TestAzureResponsesAPIConfig: == "https://litellm8397336933.openai.azure.com/openai/v1/responses?api-version=latest" ) - # Test with date-based version - should use openai/responses - result_date = self.config.get_complete_url( + result_dated_preview = self.config.get_complete_url( api_base=base_url, - litellm_params={"api_version": "2025-01-01"}, + litellm_params={"api_version": "2025-04-01-preview"}, ) assert ( - result_date - == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-01-01" + result_dated_preview + == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" ) + @pytest.mark.parametrize("api_version", ["2024-10-21", "2024-06-01", "2025-01-01"]) + def test_dated_ga_api_version_routes_responses_like_an_unset_one(self, api_version: str): + base_url: Final = "https://litellm8397336933.openai.azure.com" + + dated: Final = self.config.get_complete_url(api_base=base_url, litellm_params={"api_version": api_version}) + unset: Final = self.config.get_complete_url(api_base=base_url, litellm_params={}) + + assert dated == unset + assert httpx.URL(dated).params["api-version"] == AZURE_DEFAULT_RESPONSES_API_VERSION + def test_azure_get_complete_url_with_default_api_version(self): """Test Azure get_complete_url uses default API version when none is provided""" from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION @@ -320,6 +333,64 @@ class TestAzureResponsesAPIConfig: ) assert result_none_version == expected_url + @pytest.mark.parametrize( + ("api_version", "responses_api_version", "expected_path"), + [ + ("2024-10-21", "2025-04-01-preview", "/openai/responses?api-version=2025-04-01-preview"), + ("2025-04-01-preview", "v1", "/openai/v1/responses?api-version=v1"), + (None, "2025-04-01-preview", "/openai/responses?api-version=2025-04-01-preview"), + ], + ) + def test_responses_api_version_wins_over_api_version( + self, api_version: str | None, responses_api_version: str, expected_path: str + ): + base_url: Final = "https://litellm8397336933.openai.azure.com" + + url: Final = self.config.get_complete_url( + api_base=base_url, + litellm_params={"api_version": api_version, "responses_api_version": responses_api_version}, + ) + + assert url == f"{base_url}{expected_path}" + + @pytest.mark.parametrize("responses_api_version", [None, "v1"]) + def test_api_version_in_api_base_query_wins_over_a_dated_ga_api_version(self, responses_api_version: str | None): + api_base: Final = "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" + + url: Final = self.config.get_complete_url( + api_base=api_base, + litellm_params={"api_version": "2024-10-21", "responses_api_version": responses_api_version}, + ) + + assert url == api_base + + def test_api_version_in_env_api_base_query_wins_over_a_dated_ga_api_version(self, monkeypatch: pytest.MonkeyPatch): + env_api_base: Final = ( + "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" + ) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setenv("AZURE_API_BASE", env_api_base) + + url: Final = self.config.get_complete_url(api_base=None, litellm_params={"api_version": "2024-10-21"}) + + assert url == env_api_base + + @pytest.mark.parametrize( + ("api_base", "expected_path"), + [ + ("https://apim.example.com/openai-gw/team", "/openai-gw/team/openai/v1/responses"), + ("https://apim.example.com/openai/team-a", "/openai/team-a/openai/v1/responses"), + ], + ) + @pytest.mark.parametrize("litellm_params", [{"api_version": "2024-10-21"}, {"responses_api_version": "v1"}]) + def test_v1_route_only_rewrites_the_responses_route_segment_of_a_gateway_api_base( + self, api_base: str, expected_path: str, litellm_params: dict[str, str] + ): + url: Final = httpx.URL(self.config.get_complete_url(api_base=api_base, litellm_params=litellm_params)) + + assert url.host == "apim.example.com" + assert url.path == expected_path + def test_azure_cancel_response_api_request(self): """Test Azure cancel response API request transformation""" from litellm.types.router import GenericLiteLLMParams @@ -684,6 +755,207 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo ) +_COMPLETED_RESPONSE: Final = { + "id": "resp_dated_chat_deployment", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-5.4-nano", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, +} + + +_AZURE_RESOURCE: Final = "https://example-resource.openai.azure.com" + + +def _unset_api_version_responses_url() -> httpx.URL: + return httpx.URL(AzureOpenAIResponsesAPIConfig().get_complete_url(api_base=_AZURE_RESOURCE, litellm_params={})) + + +def _azure_responses_create(api_version: str) -> str | None: + response: Final = litellm.responses( + model="azure/gpt-5.4-nano", + input="hi", + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version=api_version, + ) + return response.output[0].content[0].text + + +def _azure_chat_bridged_to_responses(api_version: str) -> str | None: + response: Final = litellm.completion( + model="azure/responses/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version=api_version, + ) + return response.choices[0].message.content + + +@pytest.mark.parametrize("call_responses", [_azure_responses_create, _azure_chat_bridged_to_responses]) +def test_dated_ga_api_version_sends_responses_to_the_unset_api_version_url( + respx_mock: respx.MockRouter, call_responses: Callable[[str], str | None] +) -> None: + route: Final = respx_mock.post(host="example-resource.openai.azure.com").mock( + return_value=httpx.Response(200, json=_COMPLETED_RESPONSE) + ) + + text: Final = call_responses("2024-10-21") + + assert text == "Done." + assert route.call_count == 1 + assert route.calls[0].request.url == _unset_api_version_responses_url() + + +_LIFECYCLE_BODIES: Final = { + "get": _COMPLETED_RESPONSE, + "delete": {"id": "resp_dated", "object": "response", "deleted": True}, + "cancel": {**_COMPLETED_RESPONSE, "status": "cancelled"}, + "input_items": {"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}, +} + + +def _call_lifecycle(operation: str) -> object: + kwargs: Final = { + "response_id": "resp_dated", + "custom_llm_provider": "azure", + "api_base": _AZURE_RESOURCE, + "api_key": "azure-test-key", + "api_version": "2024-10-21", + } + match operation: + case "get": + return litellm.get_responses(**kwargs) + case "delete": + return litellm.delete_responses(**kwargs) + case "cancel": + return litellm.cancel_responses(**kwargs) + case _: + return litellm.list_input_items(**kwargs) + + +@pytest.mark.parametrize( + ("operation", "path_suffix"), + [("get", ""), ("delete", ""), ("cancel", "/cancel"), ("input_items", "/input_items")], +) +def test_dated_ga_api_version_sends_response_lifecycle_calls_to_the_unset_api_version_route( + respx_mock: respx.MockRouter, operation: str, path_suffix: str +) -> None: + route: Final = respx_mock.route(host="example-resource.openai.azure.com").mock( + return_value=httpx.Response(200, json=_LIFECYCLE_BODIES[operation]) + ) + + _call_lifecycle(operation) + + unset_url: Final = _unset_api_version_responses_url() + sent_url: Final = route.calls[0].request.url + assert route.call_count == 1 + assert sent_url.path == f"{unset_url.path}/resp_dated{path_suffix}" + assert sent_url.params["api-version"] == unset_url.params["api-version"] + + +def test_chat_completion_keeps_the_dated_ga_api_version(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post( + f"{_AZURE_RESOURCE}/openai/deployments/gpt-5.4-nano/chat/completions", params={"api-version": "2024-10-21"} + ).mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-dated", + "object": "chat.completion", + "created": 1734366691, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Done."}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.completion( + model="azure/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version="2024-10-21", + ) + + assert response.choices[0].message.content == "Done." + assert route.call_count == 1 + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_routes_a_dated_chat_deployment_to_its_own_responses_version( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post( + "https://example-resource.openai.azure.com/openai/responses", params={"api-version": "2025-04-01-preview"} + ).mock(return_value=httpx.Response(200, json=_COMPLETED_RESPONSE)) + + response: Final = litellm.responses( + model="azure/gpt-5.4-nano", + input="hi", + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-10-21", + responses_api_version="2025-04-01-preview", + ) + request_body: Final = json.loads(route.calls[0].request.content) + + assert route.call_count == 1 + assert response.output[0].content[0].text == "Done." + assert request_body["model"] == "gpt-5.4-nano" + assert "responses_api_version" not in request_body + assert "api_version" not in request_body + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_reaches_the_url_when_chat_completion_bridges_to_responses( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post( + "https://example-resource.openai.azure.com/openai/v1/responses", params={"api-version": "v1"} + ).mock(return_value=httpx.Response(200, json=_COMPLETED_RESPONSE)) + + response: Final = litellm.completion( + model="azure/responses/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-10-21", + responses_api_version="v1", + ) + request_body: Final = json.loads(route.calls[0].request.content) + + assert route.call_count == 1 + assert response.choices[0].message.content == "Done." + assert "responses_api_version" not in request_body + + def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap(): request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request( model="azure_ai/gpt-5.4-nano", diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 9514d64615a..a445d951bc2 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -345,6 +345,80 @@ def test_azure_ai_grok_stop_parameter_handling(): assert "stop" in gpt_params, "GPT models should support stop parameter" +@pytest.mark.respx(assert_all_called=True) +def test_azure_ai_grok_bills_reasoning_tokens_reported_outside_completion_tokens( + respx_mock: respx.MockRouter, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + route: Final = respx_mock.post(url__startswith=api_url).mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-grok", + "object": "chat.completion", + "created": 1700000000, + "model": "grok-4.7", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "42"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 14, + "completion_tokens": 10, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + }, + ) + ) + + response: Final = litellm.completion( + model="azure_ai/grok-4.7", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + ) + model_info: Final = litellm.get_model_info(model="grok-4.7", custom_llm_provider="azure_ai") + expected_cost: Final = 14 * model_info["input_cost_per_token"] + 322 * model_info["output_cost_per_token"] + + assert route.call_count == 1 + assert response.usage.completion_tokens == 322 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost, rel=1e-9) + + +@pytest.mark.respx(assert_all_called=True) +def test_azure_ai_grok_stream_folds_reasoning_tokens_into_the_usage_chunk(respx_mock: respx.MockRouter) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + stream_body: Final = ( + 'data: {"id":"chatcmpl-grok","object":"chat.completion.chunk","created":1700000000,"model":"grok-4.7",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"42"},"finish_reason":null}]}\n\n' + 'data: {"id":"chatcmpl-grok","object":"chat.completion.chunk","created":1700000000,"model":"grok-4.7",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n' + 'data: {"id":"chatcmpl-grok","object":"chat.completion.chunk","created":1700000000,"model":"grok-4.7",' + '"choices":[],"usage":{"prompt_tokens":14,"completion_tokens":10,"total_tokens":336,' + '"completion_tokens_details":{"reasoning_tokens":312}}}\n\n' + "data: [DONE]\n\n" + ) + route: Final = respx_mock.post(url__startswith=api_url).mock( + return_value=Response(200, text=stream_body, headers={"content-type": "text/event-stream"}) + ) + + chunks: Final = tuple( + litellm.completion( + model="azure_ai/grok-4.7", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + stream=True, + stream_options={"include_usage": True}, + ) + ) + usage_chunks: Final = [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + assert route.call_count == 1 + assert len(usage_chunks) == 1 + assert usage_chunks[0].completion_tokens == 322 + assert usage_chunks[0].total_tokens == usage_chunks[0].prompt_tokens + usage_chunks[0].completion_tokens + + def test_azure_model_router_response_shows_actual_model(): """ Test that Azure Model Router returns the actual model used in the response, diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 72857246d85..16f44f4f35f 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1087,6 +1087,35 @@ def test_bedrock_chat_invoke_eager_input_streaming_beta_not_duplicated_with_clie assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" + + +def _chat_invoke_betas(extra_optional_params: dict[str, object]) -> object: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params: Final = config.map_openai_params( + non_default_params={"max_tokens": 64}, optional_params={}, model=model, drop_params=False + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={**optional_params, **extra_optional_params}, + litellm_params={}, + headers={"anthropic-beta": f"claude-code-20250219,{DANGEROUS_TOOL_USE_BETA}"}, + ).get("anthropic_beta") + + +def test_bedrock_chat_invoke_drops_client_dangerous_tool_use_beta_without_safeguards() -> None: + assert _chat_invoke_betas({}) == ["claude-code-20250219"] + + +def test_bedrock_chat_invoke_keeps_client_dangerous_tool_use_beta_with_safeguards() -> None: + assert _chat_invoke_betas({"safeguards": [{"type": "dangerous_tool_use"}]}) == [ + "claude-code-20250219", + DANGEROUS_TOOL_USE_BETA, + ] + + def _mid_conversation_system_conversation() -> list[dict]: return [ { diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index dffe2c36e20..c91eb75dbd9 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1677,6 +1677,60 @@ def test_bedrock_messages_does_not_add_dangerous_tool_use_beta_without_safeguard assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) +_NAME_SCHEMA: Final = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]} + + +@pytest.mark.parametrize( + "request_params", + [ + {}, + {"output_config": {"format": {"type": "json_schema", "schema": _NAME_SCHEMA}}}, + {"tools": [{"name": "get_name", "strict": True, "input_schema": _NAME_SCHEMA}]}, + ], + ids=["plain", "output_config_format", "strict_tool"], +) +def test_bedrock_messages_drops_client_dangerous_tool_use_beta_without_safeguards( + local_model_cost_map: None, local_beta_headers_config: None, request_params: dict[str, object] +) -> None: + from litellm.types.router import GenericLiteLLMParams + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": [{"type": "text", "text": "Name a color."}]}], + anthropic_messages_optional_request_params={"max_tokens": 64, **request_params}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}, + ) + + assert all(result.get(key) == value for key, value in request_params.items()) + assert "dangerous-tool-use-2026-09-03" not in result.get("anthropic_beta", []) + + +def test_bedrock_messages_keeps_client_dangerous_tool_use_beta_with_safeguards_and_structured_outputs( + local_model_cost_map: None, local_beta_headers_config: None +) -> None: + from litellm.types.router import GenericLiteLLMParams + + safeguards: Final = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + strict_tool: Final = {"name": "get_name", "strict": True, "input_schema": _NAME_SCHEMA} + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Name a color."}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 64, + "safeguards": safeguards, + "tools": [strict_tool], + }, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}, + ) + + assert result["safeguards"] == safeguards + assert result["tools"] == [strict_tool] + assert result["anthropic_beta"].count("dangerous-tool-use-2026-09-03") == 1 + + def test_bedrock_messages_stream_decoder_keeps_safeguard_results(): """Bedrock streams the classifier verdicts on message_start and on the final message_delta, exactly as api.anthropic.com does.""" decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="us.anthropic.claude-sonnet-5") diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py index 923572c4f46..3236efbbfba 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_messages_transformation.py @@ -10,6 +10,7 @@ Bearer/SigV4 auth chain, and the wire request through the public entrypoint. """ import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -453,6 +454,14 @@ class TestBetaHeadersOnTheWire: assert _sent_betas(route) == ["dangerous-tool-use-2026-09-03"] assert _sent_body(route)["safeguards"] == safeguards + @pytest.mark.asyncio + @respx.mock + async def test_client_dangerous_tool_use_beta_reaches_mantle_without_safeguards(self) -> None: + route: Final = await self._send(extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03"}) + + assert _sent_betas(route) == ["dangerous-tool-use-2026-09-03"] + assert "safeguards" not in _sent_body(route) + @pytest.mark.asyncio @respx.mock async def test_betas_and_version_never_travel_in_the_body(self): diff --git a/tests/unit/llms/openai/completion/test_handler.py b/tests/unit/llms/openai/completion/test_handler.py index b4b224f0977..6a7302ec122 100644 --- a/tests/unit/llms/openai/completion/test_handler.py +++ b/tests/unit/llms/openai/completion/test_handler.py @@ -106,7 +106,7 @@ async def test_acompletion_uses_optimized_http_client(): with patch.object(BaseOpenAILLM, "_get_async_http_client", return_value=mock_http_client) as mock_get_client: with patch( - "litellm.llms.openai.completion.handler.AsyncOpenAI", + "litellm.llms.openai.workload_identity.AsyncOpenAI", return_value=mock_async_openai, ) as mock_openai_class: handler = OpenAITextCompletion() diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 0da28e7207e..f85e4db447b 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -1434,11 +1434,11 @@ class TestAzureResponsesAPIConfig: # Test with date-based version - should use openai/responses result_date = self.config.get_complete_url( api_base=base_url, - litellm_params={"api_version": "2025-01-01"}, + litellm_params={"api_version": "2025-04-01-preview"}, ) assert ( result_date - == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-01-01" + == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" ) def test_azure_transform_then_normalize_strips_custom_tool_call_namespace(self): diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 3c7dff34296..831df30c98a 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -442,15 +442,26 @@ def _sdk_api_client( ) if api == "image_variations": variation_factory: Final = OpenAIImageVariationsHandler() - params: Final = { - "api_key": "transport-only", - "base_url": kwargs["api_base"], - "timeout": timeout, - } return ( - variation_factory.get_async_client(client=client, init_client_params=params) + variation_factory.get_async_client( + client=client, + api_key="transport-only", + api_base="https://sdk-default.example/v1", + timeout=timeout, + max_retries=0, + organization=None, + litellm_params={}, + ) if is_async - else variation_factory.get_sync_client(client=client, init_client_params=params) + else variation_factory.get_sync_client( + client=client, + api_key="transport-only", + api_base="https://sdk-default.example/v1", + timeout=timeout, + max_retries=0, + organization=None, + litellm_params={}, + ) ) if api == "azure_gateway": return BaseAzureLLM()._init_azure_client_for_cloudflare_ai_gateway( diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index 2236de1eb5f..1a1a4ca7047 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -1,5 +1,9 @@ +import asyncio +import io import json import sys +import threading +from collections.abc import Awaitable, Callable from functools import partial from pathlib import Path from typing import Final @@ -7,21 +11,36 @@ from typing import Final import httpx import pytest import respx -from openai import AsyncOpenAI, OpenAI +from openai import DEFAULT_MAX_RETRIES as OPENAI_SDK_DEFAULT_MAX_RETRIES +from openai import APIStatusError, AsyncOpenAI, OpenAI import litellm +from litellm.containers import alist_container_files, aretrieve_container, list_container_files +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError +from litellm.llms.openai.containers.transformation import OpenAIContainerConfig +from litellm.llms.openai.evals.transformation import OpenAIEvalsConfig +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.openai.vector_store_files.transformation import OpenAIVectorStoreFilesConfig +from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.llms.openai.workload_identity import ( OpenAIWorkloadIdentityConfig, _workload_identity_auth, + build_async_openai_client, + build_openai_client, get_workload_identity_bearer_token, + resolve_openai_bearer_token, resolve_openai_workload_identity_config, ) from litellm.types.router import GenericLiteLLMParams +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.workload_identity import OPENAI_WIF_KWARGS_KEYS TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token" CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions" @@ -190,6 +209,25 @@ class TestTokenExchange: assert first == second == "exchanged-bearer-token" assert route.call_count == 1 + @respx.mock + def test_more_identities_than_the_old_cache_ceiling_stay_cached( + self, wif_env: OpenAIWorkloadIdentityConfig + ) -> None: + route: Final = mock_token_exchange() + identities: Final = tuple( + OpenAIWorkloadIdentityConfig( + identity_provider_id=wif_env.identity_provider_id, + service_account_id=f"user-{index}", + token_file=wif_env.token_file, + ) + for index in range(24) + ) + for identity in identities: + assert get_workload_identity_bearer_token(identity) == "exchanged-bearer-token" + for identity in identities: + assert get_workload_identity_bearer_token(identity) == "exchanged-bearer-token" + assert route.call_count == len(identities) + def test_old_sdk_raises_upgrade_error( self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -410,9 +448,60 @@ class TestResolveConfigFromDeployment: ) assert config == wif_env - def test_partial_litellm_params_without_env_disable(self, deployment_wif: dict[str, str]) -> None: - partial: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"} - assert resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=partial) is None + @pytest.mark.parametrize("missing_key", sorted(OPENAI_WIF_KWARGS_KEYS)) + def test_partial_litellm_params_without_env_raise_naming_the_missing_field( + self, deployment_wif: dict[str, str], missing_key: str + ) -> None: + subset: Final = {key: value for key, value in deployment_wif.items() if key != missing_key} + with pytest.raises(OpenAIError) as error: + resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=subset) + assert error.value.status_code == 401 + assert f"missing {missing_key}." in error.value.message + assert all(key not in error.value.message.split(".")[0] for key in subset) + + @respx.mock + @pytest.mark.parametrize( + "call", + [ + partial(litellm.completion, model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}]), + partial(litellm.image_generation, model="openai/gpt-image-2", prompt="a cat"), + ], + ids=["chat", "image"], + ) + def test_partial_deployment_identity_surfaces_as_an_auth_error_naming_the_gap( + self, deployment_wif: dict[str, str], call: Callable[..., object] + ) -> None: + subset: Final = {key: value for key, value in deployment_wif.items() if key != "openai_service_account_id"} + with pytest.raises(litellm.AuthenticationError) as error: + call(**subset) + assert error.value.status_code == 401 + assert "missing openai_service_account_id." in error.value.message + + def test_partial_litellm_params_fall_back_to_the_process_key( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + subset: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"} + assert ( + resolve_openai_workload_identity_config(api_key="sk-from-env", api_base=None, litellm_params=subset) is None + ) + + def test_deployment_key_equal_to_the_process_key_still_wins( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-shared") + assert ( + resolve_openai_workload_identity_config( + api_key="sk-shared", api_base=None, litellm_params={**deployment_wif, "api_key": "sk-shared"} + ) + is None + ) + + def test_partial_litellm_params_yield_to_the_deployments_own_key(self, deployment_wif: dict[str, str]) -> None: + subset: Final = {key: value for key, value in deployment_wif.items() if key != "openai_identity_token_file"} + assert ( + resolve_openai_workload_identity_config(api_key="sk-static", api_base=None, litellm_params=subset) is None + ) def test_static_api_key_beats_litellm_params(self, deployment_wif: dict[str, str]) -> None: assert ( @@ -420,7 +509,7 @@ class TestResolveConfigFromDeployment: is None ) - def test_env_openai_api_key_beats_litellm_params( + def test_env_openai_api_key_beats_deployment_identity( self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") @@ -428,6 +517,16 @@ class TestResolveConfigFromDeployment: resolve_openai_workload_identity_config(api_key=None, api_base=None, litellm_params=deployment_wif) is None ) + @pytest.mark.parametrize("global_attr", ["api_key", "openai_key"]) + def test_litellm_module_key_beats_deployment_identity( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch, global_attr: str + ) -> None: + monkeypatch.setattr(litellm, global_attr, "sk-module-global") + assert ( + resolve_openai_bearer_token(api_key=None, api_base=None, litellm_params=deployment_wif) + == "sk-module-global" + ) + def test_foreign_api_base_disables_deployment_wif(self, deployment_wif: dict[str, str]) -> None: assert ( resolve_openai_workload_identity_config( @@ -468,6 +567,108 @@ class TestDeploymentClientConstruction: assert first is not second assert again is first + def test_builders_reuse_cached_client_per_identity(self, deployment_wif: dict[str, str]) -> None: + other_deployment: Final = {**deployment_wif, "openai_service_account_id": "user-other"} + first: Final = build_openai_client(api_key=None, api_base=None, litellm_params=deployment_wif) + again: Final = build_openai_client(api_key=None, api_base=None, litellm_params=dict(deployment_wif)) + other: Final = build_openai_client(api_key=None, api_base=None, litellm_params=other_deployment) + async_first: Final = build_async_openai_client(api_key=None, api_base=None, litellm_params=deployment_wif) + async_again: Final = build_async_openai_client(api_key=None, api_base=None, litellm_params=deployment_wif) + assert again is first + assert other is not first + assert async_again is async_first + assert isinstance(async_first, AsyncOpenAI) + + @pytest.mark.parametrize("builder", [build_openai_client, build_async_openai_client]) + def test_unset_max_retries_keep_the_sdk_default_and_explicit_ones_stick( + self, deployment_wif: dict[str, str], builder: Callable[..., OpenAI | AsyncOpenAI] + ) -> None: + assert ( + builder(api_key="sk-static", api_base=None, max_retries=None).max_retries == OPENAI_SDK_DEFAULT_MAX_RETRIES + ) + assert builder(api_key=None, api_base=None, max_retries=None, litellm_params=deployment_wif).max_retries == ( + OPENAI_SDK_DEFAULT_MAX_RETRIES + ) + assert builder(api_key=None, api_base=None, max_retries=0, litellm_params=deployment_wif).max_retries == 0 + + def test_builders_never_cache_static_key_clients(self) -> None: + first: Final = build_openai_client(api_key="sk-static", api_base=None, litellm_params=None) + again: Final = build_openai_client(api_key="sk-static", api_base=None, litellm_params=None) + assert again is not first + + @respx.mock + def test_shared_workload_identity_client_never_borrows_the_callers_transport( + self, deployment_wif: dict[str, str] + ) -> None: + mock_token_exchange("process-transport-bearer") + models_route: Final = respx.get(MODELS_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": []}) + ) + caller_transport: Final = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(418))) + + federated: Final = build_openai_client( + api_key=None, + api_base=None, + litellm_params=deployment_wif, + static_key_http_client_factory=lambda: caller_transport, + ) + federated.models.list() + static: Final = build_openai_client( + api_key="sk-static", + api_base=None, + litellm_params=None, + static_key_http_client_factory=lambda: caller_transport, + ) + with pytest.raises(APIStatusError) as static_error: + static.models.list() + + assert models_route.calls.last.request.headers["Authorization"] == "Bearer process-transport-bearer" + assert static_error.value.status_code == 418 + + def test_federated_builders_never_call_the_static_key_client_factory(self, deployment_wif: dict[str, str]) -> None: + def factory_that_must_stay_cold() -> httpx.Client | None: + raise AssertionError("static-key http client built on the federated path") + + async def async_factory_that_must_stay_cold() -> httpx.AsyncClient | None: + raise AssertionError("static-key async http client built on the federated path") + + assert isinstance( + build_openai_client( + api_key=None, + api_base=None, + litellm_params=deployment_wif, + static_key_http_client_factory=factory_that_must_stay_cold, + ), + OpenAI, + ) + assert isinstance( + build_async_openai_client( + api_key=None, + api_base=None, + litellm_params=deployment_wif, + static_key_http_client_factory=async_factory_that_must_stay_cold, + ), + AsyncOpenAI, + ) + + @respx.mock + @pytest.mark.asyncio + async def test_shared_workload_identity_client_rides_the_process_session( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + mock_token_exchange("session-bearer") + session_models: Final = {"object": "list", "data": [{"id": "model-from-process-session", "object": "model"}]} + monkeypatch.setattr( + litellm, + "aclient_session", + httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, json=session_models))), + ) + + federated: Final = build_async_openai_client(api_key=None, api_base=None, litellm_params=deployment_wif) + listed: Final = await federated.models.list() + + assert [model.id for model in listed.data] == ["model-from-process-session"] + @respx.mock def test_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: mock_token_exchange("deployment-bearer") @@ -565,6 +766,145 @@ class TestResponsesValidateEnvironmentFromDeployment: ) assert headers["Authorization"] == "Bearer sk-responses" + def test_static_key_equal_to_the_process_key_wins( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-shared") + assert ( + OpenAIResponsesAPIConfig().validate_environment( + headers={}, + model="gpt-4o-mini", + litellm_params=GenericLiteLLMParams(api_key="sk-shared", **deployment_wif), + )["Authorization"] + == "Bearer sk-shared" + ) + + +class TestResolveBearerToken: + @respx.mock + def test_keyless_deployment_mints_the_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + exchange: Final = mock_token_exchange("resolved-bearer") + assert ( + resolve_openai_bearer_token(api_key=None, api_base=None, litellm_params=deployment_wif) == "resolved-bearer" + ) + assert exchange.call_count == 1 + + @respx.mock + def test_deployment_key_wins_without_an_exchange(self, deployment_wif: dict[str, str]) -> None: + exchange: Final = mock_token_exchange() + assert ( + resolve_openai_bearer_token(api_key="sk-static", api_base=None, litellm_params=deployment_wif) + == "sk-static" + ) + assert exchange.call_count == 0 + + @respx.mock + def test_process_wide_key_beats_deployment_identity_without_an_exchange( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + exchange: Final = mock_token_exchange() + assert resolve_openai_bearer_token(api_key=None, api_base=None, litellm_params=deployment_wif) == "sk-from-env" + assert exchange.call_count == 0 + + def test_process_wide_key_serves_deployments_without_identity(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + assert ( + resolve_openai_bearer_token(api_key=None, api_base=None, litellm_params={"model": "gpt-5.6"}) + == "sk-from-env" + ) + + def test_nothing_configured_resolves_to_none(self, deployment_wif: dict[str, str]) -> None: + assert resolve_openai_bearer_token(api_key=None, api_base=None, litellm_params=None) is None + + +HeadersFromParams = Callable[[GenericLiteLLMParams], dict[str, str]] +HTTP_HANDLER_VALIDATORS: Final[tuple[tuple[str, HeadersFromParams], ...]] = ( + ( + "image_edit", + lambda params: OpenAIImageEditConfig().validate_environment( + headers={}, + model="gpt-image-2", + api_key=params.api_key, + litellm_params=dict(params), + api_base=params.api_base, + ), + ), + ("vector_stores", lambda params: OpenAIVectorStoreConfig().validate_environment(headers={}, litellm_params=params)), + ( + "vector_store_files", + lambda params: OpenAIVectorStoreFilesConfig().validate_environment(headers={}, litellm_params=params), + ), + ( + "videos", + lambda params: OpenAIVideoConfig().validate_environment( + headers={}, model="sora-2", api_key=params.api_key, litellm_params=params + ), + ), + ( + "containers", + lambda params: OpenAIContainerConfig().validate_environment( + headers={}, api_key=params.api_key, litellm_params=params + ), + ), + ("evals", lambda params: OpenAIEvalsConfig().validate_environment(headers={}, litellm_params=params)), +) + + +@pytest.mark.parametrize( + ("surface", "headers_from"), HTTP_HANDLER_VALIDATORS, ids=[name for name, _ in HTTP_HANDLER_VALIDATORS] +) +class TestHttpHandlersHonorDeploymentIdentity: + @respx.mock + def test_keyless_deployment_mints_the_exchanged_bearer( + self, deployment_wif: dict[str, str], surface: str, headers_from: HeadersFromParams + ) -> None: + exchange: Final = mock_token_exchange(f"{surface}-bearer") + headers: Final = headers_from(GenericLiteLLMParams(**deployment_wif)) + assert headers["Authorization"] == f"Bearer {surface}-bearer" + assert exchange.call_count == 1 + + @respx.mock + def test_deployment_key_wins_without_an_exchange( + self, deployment_wif: dict[str, str], surface: str, headers_from: HeadersFromParams + ) -> None: + exchange: Final = mock_token_exchange() + headers: Final = headers_from(GenericLiteLLMParams(api_key=f"sk-{surface}", **deployment_wif)) + assert headers["Authorization"] == f"Bearer sk-{surface}" + assert exchange.call_count == 0 + + @respx.mock + def test_process_wide_key_beats_deployment_identity_without_an_exchange( + self, + deployment_wif: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + surface: str, + headers_from: HeadersFromParams, + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + exchange: Final = mock_token_exchange() + headers: Final = headers_from(GenericLiteLLMParams(**deployment_wif)) + assert headers["Authorization"] == "Bearer sk-from-env" + assert exchange.call_count == 0 + + def test_process_wide_key_still_serves_deployments_without_identity( + self, + deployment_wif: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + surface: str, + headers_from: HeadersFromParams, + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + headers: Final = headers_from(GenericLiteLLMParams()) + assert headers["Authorization"] == "Bearer sk-from-env" + + def test_partial_deployment_identity_fails_loudly( + self, deployment_wif: dict[str, str], surface: str, headers_from: HeadersFromParams + ) -> None: + subset: Final = {key: value for key, value in deployment_wif.items() if key != "openai_service_account_id"} + with pytest.raises(OpenAIError, match="missing openai_service_account_id"): + headers_from(GenericLiteLLMParams(**subset)) + class TestDiscoverModels: @staticmethod @@ -650,6 +990,967 @@ class TestClientsideBaseOverride: ) +IMAGES_URL: Final = "https://api.openai.com/v1/images/generations" +IMAGE_VARIATIONS_URL: Final = "https://api.openai.com/v1/images/variations" +SPEECH_URL: Final = "https://api.openai.com/v1/audio/speech" +TRANSCRIPTIONS_URL: Final = "https://api.openai.com/v1/audio/transcriptions" +MODERATIONS_URL: Final = "https://api.openai.com/v1/moderations" +COMPLETIONS_URL: Final = "https://api.openai.com/v1/completions" +FILES_URL: Final = "https://api.openai.com/v1/files" +BATCHES_URL: Final = "https://api.openai.com/v1/batches" +FINE_TUNING_JOBS_URL: Final = "https://api.openai.com/v1/fine_tuning/jobs" +ASSISTANTS_URL: Final = "https://api.openai.com/v1/assistants" +THREADS_URL: Final = "https://api.openai.com/v1/threads" +CONTAINERS_URL: Final = "https://api.openai.com/v1/containers" +IMAGE_BODY: Final = {"created": 1, "data": [{"b64_json": "aGk="}]} +MODERATION_BODY: Final = { + "id": "modr-1", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}}], +} +TEXT_COMPLETION_BODY: Final = { + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +FILE_BODY: Final = { + "id": "file-1", + "object": "file", + "bytes": 2, + "created_at": 1, + "filename": "in.jsonl", + "purpose": "batch", + "status": "processed", +} +FILE_DELETED_BODY: Final = {"id": "file-1", "object": "file", "deleted": True} +BATCH_BODY: Final = { + "id": "batch-1", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-1", + "completion_window": "24h", + "status": "validating", + "created_at": 1, +} +FINE_TUNING_JOB_BODY: Final = { + "id": "ftjob-1", + "object": "fine_tuning.job", + "created_at": 1, + "error": None, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": "auto"}, + "model": "gpt-5.4-nano", + "organization_id": "org-1", + "result_files": [], + "seed": 0, + "status": "queued", + "trained_tokens": None, + "training_file": "file-1", + "validation_file": None, +} + + +ASSISTANT_BODY: Final = { + "id": "asst-1", + "object": "assistant", + "created_at": 1, + "name": "wif", + "description": None, + "model": "gpt-5.4-nano", + "instructions": None, + "tools": [], + "metadata": {}, +} +ASSISTANT_DELETED_BODY: Final = {"id": "asst-1", "object": "assistant.deleted", "deleted": True} +THREAD_BODY: Final = {"id": "thread-1", "object": "thread", "created_at": 1, "metadata": {}} +MESSAGE_BODY: Final = { + "id": "msg-1", + "object": "thread.message", + "created_at": 1, + "thread_id": "thread-1", + "role": "user", + "status": "completed", + "content": [{"type": "text", "text": {"value": "hi", "annotations": []}}], + "assistant_id": None, + "run_id": None, + "attachments": [], + "metadata": {}, +} +RUN_BODY: Final = { + "id": "run-1", + "object": "thread.run", + "created_at": 1, + "thread_id": "thread-1", + "assistant_id": "asst-1", + "status": "completed", + "model": "gpt-5.4-nano", + "instructions": "", + "tools": [], + "metadata": {}, + "parallel_tool_calls": False, +} + + +def mock_streaming_text_completions() -> respx.Route: + events: Final = ( + { + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "ok", "index": 0, "finish_reason": None, "logprobs": None}], + }, + { + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "", "index": 0, "finish_reason": "stop", "logprobs": None}], + }, + ) + body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n" + return respx.post(COMPLETIONS_URL).mock( + return_value=httpx.Response(200, content=body, headers={"content-type": "text/event-stream"}) + ) + + +def bearer_of(route: respx.Route) -> str: + return route.calls.last.request.headers["Authorization"] + + +class TestDeploymentNonChatSurfaces: + @respx.mock + def test_image_generation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("image-bearer") + route: Final = respx.post(IMAGES_URL).mock(return_value=httpx.Response(200, json=IMAGE_BODY)) + + response: Final = litellm.image_generation(model="openai/gpt-image-2", prompt="a cat", **deployment_wif) + + assert response.data[0].b64_json == "aGk=" + assert bearer_of(route) == "Bearer image-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_image_generation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("aimage-bearer") + route: Final = respx.post(IMAGES_URL).mock(return_value=httpx.Response(200, json=IMAGE_BODY)) + + response: Final = await litellm.aimage_generation(model="openai/gpt-image-2", prompt="a cat", **deployment_wif) + + assert response.data[0].b64_json == "aGk=" + assert bearer_of(route) == "Bearer aimage-bearer" + + @respx.mock + def test_image_variation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("variation-bearer") + route: Final = respx.post(IMAGE_VARIATIONS_URL).mock(return_value=httpx.Response(200, json=IMAGE_BODY)) + + response: Final = litellm.image_variation( + model="openai/dall-e-2", image=b"png-bytes", custom_llm_provider="openai", **deployment_wif + ) + + assert response.data[0].b64_json == "aGk=" + assert bearer_of(route) == "Bearer variation-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_image_variation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("avariation-bearer") + route: Final = respx.post(IMAGE_VARIATIONS_URL).mock(return_value=httpx.Response(200, json=IMAGE_BODY)) + + response: Final = await litellm.aimage_variation( + model="openai/dall-e-2", image=b"png-bytes", custom_llm_provider="openai", **deployment_wif + ) + + assert response.data[0].b64_json == "aGk=" + assert bearer_of(route) == "Bearer avariation-bearer" + + @respx.mock + def test_speech_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("speech-bearer") + route: Final = respx.post(SPEECH_URL).mock( + return_value=httpx.Response(200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + + response: Final = litellm.speech(model="openai/gpt-4o-mini-tts", input="hi", voice="alloy", **deployment_wif) + + assert response.content == b"audio" + assert bearer_of(route) == "Bearer speech-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_speech_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("aspeech-bearer") + route: Final = respx.post(SPEECH_URL).mock( + return_value=httpx.Response(200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + + response: Final = await litellm.aspeech( + model="openai/gpt-4o-mini-tts", input="hi", voice="alloy", **deployment_wif + ) + + assert response.content == b"audio" + assert bearer_of(route) == "Bearer aspeech-bearer" + + @respx.mock + def test_transcription_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str], tmp_path: Path) -> None: + mock_token_exchange("transcribe-bearer") + route: Final = respx.post(TRANSCRIPTIONS_URL).mock(return_value=httpx.Response(200, json={"text": "hi"})) + audio: Final = tmp_path / "tone.wav" + audio.write_bytes(b"RIFF....WAVE") + + with audio.open("rb") as audio_file: + response: Final = litellm.transcription( + model="openai/gpt-4o-mini-transcribe", file=audio_file, **deployment_wif + ) + + assert response.text == "hi" + assert bearer_of(route) == "Bearer transcribe-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_transcription_kwargs_carry_exchanged_bearer( + self, deployment_wif: dict[str, str], tmp_path: Path + ) -> None: + mock_token_exchange("atranscribe-bearer") + route: Final = respx.post(TRANSCRIPTIONS_URL).mock(return_value=httpx.Response(200, json={"text": "hi"})) + audio: Final = tmp_path / "tone.wav" + audio.write_bytes(b"RIFF....WAVE") + + with audio.open("rb") as audio_file: + response: Final = await litellm.atranscription( + model="openai/gpt-4o-mini-transcribe", file=audio_file, **deployment_wif + ) + + assert response.text == "hi" + assert bearer_of(route) == "Bearer atranscribe-bearer" + + @respx.mock + def test_moderation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("moderation-bearer") + route: Final = respx.post(MODERATIONS_URL).mock(return_value=httpx.Response(200, json=MODERATION_BODY)) + + response: Final = litellm.moderation(input="hi", model="omni-moderation-latest", **deployment_wif) + + assert response.results[0].flagged is False + assert bearer_of(route) == "Bearer moderation-bearer" + + @respx.mock + def test_moderation_ignores_unrelated_kwargs_it_never_validated(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("moderation-bearer") + route: Final = respx.post(MODERATIONS_URL).mock(return_value=httpx.Response(200, json=MODERATION_BODY)) + + response: Final = litellm.moderation( + input="hi", model="omni-moderation-latest", api_version=20240101, **deployment_wif + ) + + assert response.results[0].flagged is False + assert bearer_of(route) == "Bearer moderation-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_moderation_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("amoderation-bearer") + route: Final = respx.post(MODERATIONS_URL).mock(return_value=httpx.Response(200, json=MODERATION_BODY)) + + response: Final = await litellm.amoderation(input="hi", model="omni-moderation-latest", **deployment_wif) + + assert response.results[0].flagged is False + assert bearer_of(route) == "Bearer amoderation-bearer" + + @respx.mock + def test_text_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("text-bearer") + route: Final = respx.post(COMPLETIONS_URL).mock(return_value=httpx.Response(200, json=TEXT_COMPLETION_BODY)) + + response: Final = litellm.text_completion( + model="text-completion-openai/gpt-3.5-turbo-instruct", prompt="hi", **deployment_wif + ) + + assert response.choices[0].text == "ok" + assert bearer_of(route) == "Bearer text-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_text_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("atext-bearer") + route: Final = respx.post(COMPLETIONS_URL).mock(return_value=httpx.Response(200, json=TEXT_COMPLETION_BODY)) + + response: Final = await litellm.atext_completion( + model="text-completion-openai/gpt-3.5-turbo-instruct", prompt="hi", **deployment_wif + ) + + assert response.choices[0].text == "ok" + assert bearer_of(route) == "Bearer atext-bearer" + + @respx.mock + def test_streaming_text_completion_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("text-stream-bearer") + route: Final = mock_streaming_text_completions() + + chunks: Final = tuple( + litellm.text_completion( + model="text-completion-openai/gpt-3.5-turbo-instruct", prompt="hi", stream=True, **deployment_wif + ) + ) + + assert chunks[0].choices[0].text == "ok" + assert bearer_of(route) == "Bearer text-stream-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_streaming_text_completion_kwargs_carry_exchanged_bearer( + self, deployment_wif: dict[str, str] + ) -> None: + mock_token_exchange("atext-stream-bearer") + route: Final = mock_streaming_text_completions() + + stream: Final = await litellm.atext_completion( + model="text-completion-openai/gpt-3.5-turbo-instruct", prompt="hi", stream=True, **deployment_wif + ) + chunks: Final = [chunk async for chunk in stream] + + assert chunks[0].choices[0].text == "ok" + assert bearer_of(route) == "Bearer atext-stream-bearer" + + @respx.mock + @pytest.mark.parametrize( + ("method", "url", "response", "invoke"), + [ + pytest.param( + "POST", + FILES_URL, + httpx.Response(200, json=FILE_BODY), + lambda p: litellm.create_file(file=("in.jsonl", b"{}"), purpose="batch", **p), + id="create_file", + ), + pytest.param( + "GET", + f"{FILES_URL}/file-1", + httpx.Response(200, json=FILE_BODY), + lambda p: litellm.file_retrieve(file_id="file-1", **p), + id="file_retrieve", + ), + pytest.param( + "DELETE", + f"{FILES_URL}/file-1", + httpx.Response(200, json=FILE_DELETED_BODY), + lambda p: litellm.file_delete(file_id="file-1", **p), + id="file_delete", + ), + pytest.param( + "GET", + FILES_URL, + httpx.Response(200, json={"object": "list", "data": [FILE_BODY]}), + lambda p: litellm.file_list(**p), + id="file_list", + ), + pytest.param( + "GET", + f"{FILES_URL}/file-1/content", + httpx.Response(200, content=b"{}"), + lambda p: litellm.file_content(file_id="file-1", **p), + id="file_content", + ), + pytest.param( + "POST", + BATCHES_URL, + httpx.Response(200, json=BATCH_BODY), + lambda p: litellm.create_batch( + completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-1", **p + ), + id="create_batch", + ), + pytest.param( + "GET", + f"{BATCHES_URL}/batch-1", + httpx.Response(200, json=BATCH_BODY), + lambda p: litellm.retrieve_batch(batch_id="batch-1", **p), + id="retrieve_batch", + ), + pytest.param( + "POST", + f"{BATCHES_URL}/batch-1/cancel", + httpx.Response(200, json=BATCH_BODY), + lambda p: litellm.cancel_batch(batch_id="batch-1", **p), + id="cancel_batch", + ), + pytest.param( + "GET", + BATCHES_URL, + httpx.Response(200, json={"object": "list", "data": [BATCH_BODY], "has_more": False}), + lambda p: litellm.list_batches(**p), + id="list_batches", + ), + pytest.param( + "POST", + FINE_TUNING_JOBS_URL, + httpx.Response(200, json=FINE_TUNING_JOB_BODY), + lambda p: litellm.create_fine_tuning_job(model="gpt-5.4-nano", training_file="file-1", **p), + id="create_fine_tuning_job", + ), + pytest.param( + "POST", + f"{FINE_TUNING_JOBS_URL}/ftjob-1/cancel", + httpx.Response(200, json=FINE_TUNING_JOB_BODY), + lambda p: litellm.cancel_fine_tuning_job(fine_tuning_job_id="ftjob-1", **p), + id="cancel_fine_tuning_job", + ), + pytest.param( + "GET", + FINE_TUNING_JOBS_URL, + httpx.Response(200, json={"object": "list", "data": [FINE_TUNING_JOB_BODY], "has_more": False}), + lambda p: litellm.list_fine_tuning_jobs(**p), + id="list_fine_tuning_jobs", + ), + pytest.param( + "GET", + f"{FINE_TUNING_JOBS_URL}/ftjob-1", + httpx.Response(200, json=FINE_TUNING_JOB_BODY), + lambda p: litellm.retrieve_fine_tuning_job(fine_tuning_job_id="ftjob-1", **p), + id="retrieve_fine_tuning_job", + ), + pytest.param( + "GET", + ASSISTANTS_URL, + httpx.Response(200, json={"object": "list", "data": [ASSISTANT_BODY]}), + lambda p: litellm.get_assistants(**p), + id="get_assistants", + ), + pytest.param( + "POST", + ASSISTANTS_URL, + httpx.Response(200, json=ASSISTANT_BODY), + lambda p: litellm.create_assistants(model="gpt-5.4-nano", **p), + id="create_assistants", + ), + pytest.param( + "DELETE", + f"{ASSISTANTS_URL}/asst-1", + httpx.Response(200, json=ASSISTANT_DELETED_BODY), + lambda p: litellm.delete_assistant(assistant_id="asst-1", **p), + id="delete_assistant", + ), + pytest.param( + "POST", + THREADS_URL, + httpx.Response(200, json=THREAD_BODY), + lambda p: litellm.create_thread(**p), + id="create_thread", + ), + pytest.param( + "GET", + f"{THREADS_URL}/thread-1", + httpx.Response(200, json=THREAD_BODY), + lambda p: litellm.get_thread(thread_id="thread-1", **p), + id="get_thread", + ), + pytest.param( + "POST", + f"{THREADS_URL}/thread-1/messages", + httpx.Response(200, json=MESSAGE_BODY), + lambda p: litellm.add_message(thread_id="thread-1", role="user", content="hi", **p), + id="add_message", + ), + pytest.param( + "GET", + f"{THREADS_URL}/thread-1/messages", + httpx.Response(200, json={"object": "list", "data": [MESSAGE_BODY]}), + lambda p: litellm.get_messages(thread_id="thread-1", **p), + id="get_messages", + ), + ], + ) + def test_managed_object_kwargs_carry_exchanged_bearer( + self, + deployment_wif: dict[str, str], + method: str, + url: str, + response: httpx.Response, + invoke: Callable[[dict[str, str]], object], + ) -> None: + mock_token_exchange("managed-bearer") + route: Final = respx.route(method=method, url__startswith=url).mock(return_value=response) + + invoke({"custom_llm_provider": "openai", **deployment_wif}) + + assert bearer_of(route) == "Bearer managed-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_managed_object_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("amanaged-bearer") + files_route: Final = respx.get(url__startswith=FILES_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": [FILE_BODY]}) + ) + batches_route: Final = respx.get(url__startswith=BATCHES_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": [BATCH_BODY], "has_more": False}) + ) + jobs_route: Final = respx.get(url__startswith=FINE_TUNING_JOBS_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": [FINE_TUNING_JOB_BODY], "has_more": False}) + ) + + await litellm.afile_list(custom_llm_provider="openai", **deployment_wif) + await litellm.alist_batches(custom_llm_provider="openai", **deployment_wif) + await litellm.alist_fine_tuning_jobs(custom_llm_provider="openai", **deployment_wif) + + assert bearer_of(files_route) == "Bearer amanaged-bearer" + assert bearer_of(batches_route) == "Bearer amanaged-bearer" + assert bearer_of(jobs_route) == "Bearer amanaged-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_repeated_managed_object_calls_exchange_the_token_once(self, deployment_wif: dict[str, str]) -> None: + exchange_route: Final = mock_token_exchange("once-bearer") + files_route: Final = respx.get(url__startswith=FILES_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": [FILE_BODY]}) + ) + + for _ in range(3): + await litellm.afile_list(custom_llm_provider="openai", **deployment_wif) + + assert files_route.call_count == 3 + assert exchange_route.call_count == 1 + assert bearer_of(files_route) == "Bearer once-bearer" + + @respx.mock + def test_run_thread_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("run-bearer") + create_route: Final = respx.post(f"{THREADS_URL}/thread-1/runs").mock( + return_value=httpx.Response(200, json=RUN_BODY) + ) + poll_route: Final = respx.get(url__startswith=f"{THREADS_URL}/thread-1/runs/run-1").mock( + return_value=httpx.Response(200, json=RUN_BODY) + ) + + response: Final = litellm.run_thread( + custom_llm_provider="openai", thread_id="thread-1", assistant_id="asst-1", **deployment_wif + ) + + assert response.status == "completed" + assert bearer_of(create_route) == "Bearer run-bearer" + assert bearer_of(poll_route) == "Bearer run-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_async_assistant_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("aassistant-bearer") + assistants_route: Final = respx.get(url__startswith=ASSISTANTS_URL).mock( + return_value=httpx.Response(200, json={"object": "list", "data": [ASSISTANT_BODY]}) + ) + threads_route: Final = respx.post(url__startswith=THREADS_URL).mock( + return_value=httpx.Response(200, json=THREAD_BODY) + ) + + await litellm.aget_assistants(custom_llm_provider="openai", **deployment_wif) + await litellm.acreate_thread(custom_llm_provider="openai", **deployment_wif) + + assert bearer_of(assistants_route) == "Bearer aassistant-bearer" + assert bearer_of(threads_route) == "Bearer aassistant-bearer" + + @respx.mock + def test_container_file_kwargs_carry_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("container-files-bearer") + route: Final = respx.get(f"{CONTAINERS_URL}/cntr_1/files").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False}) + ) + + response: Final = list_container_files(container_id="cntr_1", custom_llm_provider="openai", **deployment_wif) + + assert response.has_more is False + assert bearer_of(route) == "Bearer container-files-bearer" + + @respx.mock + def test_repeated_moderation_calls_exchange_the_token_once(self, deployment_wif: dict[str, str]) -> None: + exchange_route: Final = mock_token_exchange("once-bearer") + moderation_route: Final = respx.post(MODERATIONS_URL).mock( + return_value=httpx.Response(200, json=MODERATION_BODY) + ) + + for _ in range(3): + litellm.moderation(input="hi", model="omni-moderation-latest", **deployment_wif) + + assert moderation_route.call_count == 3 + assert exchange_route.call_count == 1 + + +EMBEDDING_BODY: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, +} +DECISIONS_URL: Final = "https://api.openai.com/v1/decisions" +DECISIONS_BODY: Final = { + "model": "gpt-6-luna", + "answers": [{"type": "predicate", "name": "is_defect", "probability": 0.9}], + "usage": { + "input_tokens": 3, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 1, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 4, + }, +} +INPUT_TOKENS_URL: Final = "https://api.openai.com/v1/responses/input_tokens" +DEPLOYMENT_KEY_ROUTES: Final = ( + pytest.param( + CHAT_COMPLETIONS_URL, + httpx.Response(200, json=CHAT_COMPLETION_BODY), + lambda p: litellm.completion(model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], **p), + id="completion", + ), + pytest.param( + EMBEDDINGS_URL, + httpx.Response(200, json=EMBEDDING_BODY), + lambda p: litellm.embedding(model="openai/text-embedding-3-small", input=["hi"], **p), + id="embedding", + ), + pytest.param( + IMAGES_URL, + httpx.Response(200, json=IMAGE_BODY), + lambda p: litellm.image_generation(model="openai/gpt-image-2", prompt="a cat", **p), + id="image_generation", + ), + pytest.param( + TRANSCRIPTIONS_URL, + httpx.Response(200, json={"text": "hi"}), + lambda p: litellm.transcription(model="openai/gpt-4o-mini-transcribe", file=("tone.wav", b"RIFF"), **p), + id="transcription", + ), + pytest.param( + MODERATIONS_URL, + httpx.Response(200, json=MODERATION_BODY), + lambda p: litellm.moderation(input="hi", model="omni-moderation-latest", **p), + id="moderation", + ), + pytest.param( + MODERATIONS_URL, + httpx.Response(200, json=MODERATION_BODY), + lambda p: asyncio.run(litellm.amoderation(input="hi", model="omni-moderation-latest", **p)), + id="amoderation", + ), + pytest.param( + ASSISTANTS_URL, + httpx.Response(200, json={"object": "list", "data": [ASSISTANT_BODY]}), + lambda p: litellm.get_assistants(custom_llm_provider="openai", **p), + id="get_assistants", + ), + pytest.param( + ASSISTANTS_URL, + httpx.Response(200, json=ASSISTANT_BODY), + lambda p: litellm.create_assistants(custom_llm_provider="openai", model="gpt-5.4-nano", **p), + id="create_assistants", + ), + pytest.param( + f"{ASSISTANTS_URL}/asst-1", + httpx.Response(200, json=ASSISTANT_DELETED_BODY), + lambda p: litellm.delete_assistant(custom_llm_provider="openai", assistant_id="asst-1", **p), + id="delete_assistant", + ), + pytest.param( + DECISIONS_URL, + httpx.Response(200, json=DECISIONS_BODY), + lambda p: litellm.decisions( + model="openai/gpt-6-luna", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + **p, + ), + id="decisions", + ), +) + + +@pytest.mark.parametrize(("url", "response", "invoke"), DEPLOYMENT_KEY_ROUTES) +class TestDeploymentKeyBesideIdentity: + @respx.mock + def test_deployment_key_equal_to_the_process_key_wins_without_an_exchange( + self, + deployment_wif: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + url: str, + response: httpx.Response, + invoke: Callable[[dict[str, str]], object], + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-shared") + exchange: Final = mock_token_exchange() + route: Final = respx.route(url__startswith=url).mock(return_value=response) + + invoke({"api_key": "sk-shared", **deployment_wif}) + + assert bearer_of(route) == "Bearer sk-shared" + assert exchange.call_count == 0 + + @respx.mock + def test_process_key_beats_deployment_identity_without_an_exchange( + self, + deployment_wif: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + url: str, + response: httpx.Response, + invoke: Callable[[dict[str, str]], object], + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-shared") + exchange: Final = mock_token_exchange() + route: Final = respx.route(url__startswith=url).mock(return_value=response) + + invoke(deployment_wif) + + assert bearer_of(route) == "Bearer sk-shared" + assert exchange.call_count == 0 + + @respx.mock + def test_secret_manager_key_beats_deployment_identity_without_an_exchange( + self, + deployment_wif: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + url: str, + response: httpx.Response, + invoke: Callable[[dict[str, str]], object], + ) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", _VaultHoldingOpenAIKey()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + exchange: Final = mock_token_exchange() + route: Final = respx.route(url__startswith=url).mock(return_value=response) + + invoke(deployment_wif) + + assert bearer_of(route) == "Bearer sk-from-vault" + assert exchange.call_count == 0 + + +class _VaultHoldingOpenAIKey(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="openai_key_vault") + + async def async_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return self.sync_read_secret(secret_name, optional_params, timeout) + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return "sk-from-vault" if secret_name == "OPENAI_API_KEY" else None + + +class TestCountTokensFromDeployment: + @staticmethod + async def count(litellm_params: dict[str, str]) -> int | None: + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": litellm_params}, + ) + return None if result is None else result.total_tokens + + @respx.mock + @pytest.mark.asyncio + async def test_keyless_deployment_counts_with_the_exchanged_bearer(self, deployment_wif: dict[str, str]) -> None: + mock_token_exchange("count-bearer") + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + assert await self.count(deployment_wif) == 7 + assert bearer_of(route) == "Bearer count-bearer" + + @respx.mock + @pytest.mark.asyncio + async def test_process_key_beats_deployment_identity_without_an_exchange( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + exchange: Final = mock_token_exchange() + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + await self.count(deployment_wif) + + assert bearer_of(route) == "Bearer sk-from-env" + assert exchange.call_count == 0 + + @respx.mock + @pytest.mark.asyncio + @pytest.mark.parametrize("module_key", ["api_key", "openai_key"]) + async def test_module_level_key_beats_deployment_identity_without_an_exchange( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch, module_key: str + ) -> None: + monkeypatch.setattr(litellm, module_key, "sk-module-level") + exchange: Final = mock_token_exchange() + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + await self.count(deployment_wif) + + assert bearer_of(route) == "Bearer sk-module-level" + assert exchange.call_count == 0 + + @respx.mock + @pytest.mark.asyncio + async def test_deployment_key_wins_without_an_exchange(self, deployment_wif: dict[str, str]) -> None: + exchange: Final = mock_token_exchange() + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + await self.count({"api_key": "sk-count", **deployment_wif}) + + assert bearer_of(route) == "Bearer sk-count" + assert exchange.call_count == 0 + + @respx.mock + @pytest.mark.asyncio + @pytest.mark.parametrize( + "identity_keys", + [tuple(OPENAI_WIF_KWARGS_KEYS), ("openai_identity_provider_id",)], + ids=["complete_trio", "partial_trio"], + ) + async def test_secret_manager_key_beats_deployment_identity_without_an_exchange( + self, deployment_wif: dict[str, str], monkeypatch: pytest.MonkeyPatch, identity_keys: tuple[str, ...] + ) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", _VaultHoldingOpenAIKey()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + exchange: Final = mock_token_exchange() + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + assert await self.count({key: deployment_wif[key] for key in identity_keys}) == 7 + assert bearer_of(route) == "Bearer sk-from-vault" + assert exchange.call_count == 0 + + @respx.mock + @pytest.mark.asyncio + async def test_partial_deployment_identity_answers_a_401_naming_the_gap_without_a_count_call(self) -> None: + exchange: Final = mock_token_exchange() + route: Final = respx.post(INPUT_TOKENS_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 7})) + + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"openai_identity_provider_id": "idp-only"}}, + ) + + assert result is not None + assert (result.error, result.status_code, result.total_tokens) == (True, 401, 0) + assert "openai_service_account_id" in (result.error_message or "") + assert route.call_count == 0 + assert exchange.call_count == 0 + + +VECTOR_STORES_URL: Final = "https://api.openai.com/v1/vector_stores" +LIST_PAGE_BODY: Final = {"object": "list", "data": [], "has_more": False} + + +def mock_token_exchange_naming_its_thread(event_loop_thread: int) -> respx.Route: + def respond(request: httpx.Request) -> httpx.Response: + minted_on: Final = "event-loop" if threading.get_ident() == event_loop_thread else "worker-thread" + return httpx.Response(200, json={"access_token": f"minted-on-{minted_on}", "expires_in": 3600}) + + return respx.post(TOKEN_EXCHANGE_URL).mock(side_effect=respond) + + +LOOP_BOUND_ROUTES: Final = ( + pytest.param( + "POST", + "https://api.openai.com/v1/images/edits", + httpx.Response(200, json=IMAGE_BODY), + lambda p: litellm.aimage_edit( + model="openai/gpt-image-2", image=[io.BytesIO(b"png-bytes")], prompt="add a hat", **p + ), + id="image_edit", + ), + pytest.param( + "GET", + "https://api.openai.com/v1/videos/video_1", + httpx.Response( + 200, json={"id": "video_1", "object": "video", "status": "completed", "created_at": 1, "model": "sora-2"} + ), + lambda p: litellm.avideo_status(video_id="video_1", custom_llm_provider="openai", **p), + id="video_status", + ), + pytest.param( + "GET", + f"{CONTAINERS_URL}/cntr_1", + httpx.Response( + 200, + json={ + "id": "cntr_1", + "object": "container", + "created_at": 1, + "status": "running", + "expires_after": {"anchor": "last_active_at", "minutes": 20}, + "last_active_at": 1, + "name": "sandbox", + }, + ), + lambda p: aretrieve_container(container_id="cntr_1", custom_llm_provider="openai", **p), + id="retrieve_container", + ), + pytest.param( + "GET", + f"{CONTAINERS_URL}/cntr_1/files", + httpx.Response(200, json=LIST_PAGE_BODY), + lambda p: alist_container_files(container_id="cntr_1", custom_llm_provider="openai", **p), + id="container_files", + ), + pytest.param( + "POST", + f"{VECTOR_STORES_URL}/vs_1/search", + httpx.Response(200, json={"object": "vector_store.search_results.page", "search_query": "q", "data": []}), + lambda p: litellm.vector_stores.asearch(vector_store_id="vs_1", query="q", custom_llm_provider="openai", **p), + id="vector_store_search", + ), + pytest.param( + "GET", + f"{VECTOR_STORES_URL}/vs_1/files", + httpx.Response(200, json=LIST_PAGE_BODY), + lambda p: litellm.vector_store_files.alist(vector_store_id="vs_1", custom_llm_provider="openai", **p), + id="vector_store_files", + ), + pytest.param( + "POST", + DECISIONS_URL, + httpx.Response(200, json=DECISIONS_BODY), + lambda p: litellm.adecisions( + model="openai/gpt-6-luna", + state="review", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + **p, + ), + id="decisions", + ), +) + + +@pytest.mark.parametrize(("method", "url", "response", "invoke"), LOOP_BOUND_ROUTES) +class TestTokenExchangeLeavesTheEventLoopFree: + @respx.mock + @pytest.mark.asyncio + async def test_token_exchange_runs_off_the_event_loop( + self, + deployment_wif: dict[str, str], + method: str, + url: str, + response: httpx.Response, + invoke: Callable[[dict[str, str]], Awaitable[object]], + ) -> None: + mock_token_exchange_naming_its_thread(threading.get_ident()) + upstream: Final = respx.route(method=method, url=url).mock(return_value=response) + + await invoke(deployment_wif) + + assert bearer_of(upstream) == "Bearer minted-on-worker-thread" + + class TestCountTokens: @pytest.mark.asyncio @respx.mock diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py index a5d42e58b99..a97ab04c506 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -1,4 +1,5 @@ import asyncio +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, Mock @@ -19,6 +20,21 @@ from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +def allow_one_request(monkeypatch: pytest.MonkeyPatch, user_id: str, server_id: str) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + windows: Final = {(user_id, server_id): asyncio.Semaphore(1)} + + async def enforce(user: UserAPIKeyAuth | None, server: MCPServer) -> None: + window: Final = windows[(user.user_id if user is not None else None, server.server_id)] + if window.locked(): + raise ProxyRateLimitError(detail=f"{server_id} RPM exceeded") + await window.acquire() + + monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(enforce_mcp_server_rate_limits=enforce)) + + def context(session: ServerSession, revision: str = "2026-07-28") -> ServerRequestContext: return ServerRequestContext( session=session, @@ -356,3 +372,117 @@ async def test_proxy_route_rejects_legacy_resource_subscription_before_setup(met await getattr(server, method)(context(Mock(spec=ServerSession)), params) finally: server._mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("client_requests", "tools_list_admitted"), + [((), True), (("subscriptions/listen",), False), (("resources/subscribe",), False)], +) +async def test_only_client_initiated_listeners_consume_server_rate_limit( + monkeypatch: pytest.MonkeyPatch, client_requests: tuple[str, ...], tools_list_admitted: bool +) -> None: + from litellm.proxy._experimental.mcp_server import operations + + allow_one_request(monkeypatch, "alice", "upstream") + opened: Final = asyncio.Queue[None]() + + async def hold_upstream(*_: object, **__: object) -> None: + opened.put_nowait(None) + await asyncio.Event().wait() + + client: Final = Mock(spec=MCPClient) + client.run_with_session = AsyncMock(side_effect=hold_upstream) + server: Final = MCPServer(server_id="upstream", name="upstream", url="http://localhost/mcp", transport="http") + target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") + + async def resolve(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + return (target,) + + service: Final = subscriptions.GatewaySubscriptions(resolve_targets=resolve) + caller: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="alice")) + authorize: Final = AsyncMock(return_value=caller) + session: Final = Mock(spec=ServerSession) + session.client_params = Mock() + legacy: Final = context(session, "2025-11-25") + requests: Final = { + "subscriptions/listen": lambda: service.listen( + context(Mock(spec=ServerSession)), + SubscriptionsListenRequestParams(notifications=SubscriptionFilter(tools_list_changed=True)), + authorize, + "alice", + ), + "resources/subscribe": lambda: service.legacy_resource(legacy, "test://value", subscribe=True), + } + authorization: Final = subscriptions.SubscriptionAuthorization(caller, authorize, lambda *_: "alice") + starts: Final = (lambda: service.legacy_session(legacy, authorization), *(requests[r] for r in client_requests)) + tasks: Final = tuple(asyncio.create_task(start()) for start in starts) + try: + for _ in tasks: + await asyncio.wait_for(opened.get(), 1) + rejection: Final = await operations._mcp_server_rate_limit_rejection(server, caller.user_api_key_auth) + assert (rejection is None) is tools_list_admitted + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_legacy_listener_catalog_refreshes_do_not_consume_server_rate_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import ListToolsResult, ServerCapabilities, Tool, ToolsCapability + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog, operations + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk + + server: Final = MCPServer(server_id="list-changed", name="list-changed", transport=MCPTransport.http) + caller: Final = UserAPIKeyAuth(api_key="list-changed-key", user_id="alice") + allow_one_request(monkeypatch, "alice", "list-changed") + listed: Final = asyncio.Queue[None]() + + async def fetch_tools(_: MCPServer, **__: object) -> tuple[ListToolsResult, ServerListOk]: + listed.put_nowait(None) + return ListToolsResult(tools=[Tool(name="add", inputSchema={"type": "object"})]), ServerListOk(tool_count=1) + + monkeypatch.setenv("LITELLM_SALT_KEY", "list-changed-test-key") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(operations.global_mcp_server_manager, "registry", {server.server_id: server}) + monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])) + monkeypatch.setattr(catalog, "get_filtered_server_tools", fetch_tools) + + upstream: Final = Mock() + upstream.server_capabilities = ServerCapabilities(tools=ToolsCapability(list_changed=True)) + upstream.protocol_version = "2025-11-25" + upstream.send_ping = AsyncMock() + + async def run_with_session(connected, **_: object) -> None: + await connected(upstream) + + client: Final = Mock(spec=MCPClient) + client.run_with_session = AsyncMock(side_effect=run_with_session) + target: Final = subscriptions.SubscriptionTarget(server, client, "synthetic-fingerprint") + + async def resolve(_: OperationContext) -> tuple[subscriptions.SubscriptionTarget, ...]: + return (target,) + + initial: Final = operations.prepare_context(user_api_key_auth=caller, mcp_servers=[server.server_id]) + authorize: Final = AsyncMock(return_value=initial) + session: Final = Mock(spec=ServerSession) + session.client_params = Mock() + service: Final = subscriptions.GatewaySubscriptions(resolve_targets=resolve) + listener: Final = asyncio.create_task( + service.legacy_session( + context(session, "2025-11-25"), + subscriptions.SubscriptionAuthorization(initial, authorize, lambda *_: "alice"), + ) + ) + try: + for _ in range(2): + await asyncio.wait_for(listed.get(), 5) + assert await operations._mcp_server_rate_limit_rejection(server, caller) is None + finally: + listener.cancel() + await asyncio.gather(listener, return_exceptions=True) diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index e4129ebbc39..c7bb05e5d30 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,13 +5,16 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from litellm import Router from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.types.utils import TokenCountResponse class TestAnthropicEndpoints(unittest.TestCase): @@ -467,6 +470,151 @@ class TestErrorLogCarriesCallId: assert call_id in record.getMessage() +def _claude_router() -> Router: + return Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "sk-test"}, + } + ] + ) + + +_MALFORMED_TOOL_ENTRY: Final = [{"name": "get_weather", "input_schema": "object"}] +_SEVEN_TOKENS: Final = TokenCountResponse( + total_tokens=7, request_model="claude-sonnet", model_used="claude-sonnet", tokenizer_type="test" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "tools"), + [ + ("claude-sonnet", 5), + ("claude-sonnet", ""), + ("claude-sonnet", ["get_weather"]), + ("claude-sonnet", _MALFORMED_TOOL_ENTRY), + ("team-public-name", 5), + ], + ids=["int", "empty_string", "list_of_strings", "input_schema_string", "unnamed_model_int"], +) +async def test_count_tokens_answers_400_for_malformed_tools_without_counting(model: str, tools: object) -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock() + body: Final = {"model": model, "messages": [{"role": "user", "content": "hi"}], "tools": tools} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts it is never reached + patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response.status_code == 400 + assert json.loads(response.body)["error"]["type"] == "invalid_request_error" + counter.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_count_tokens_counts_malformed_tool_entries_for_a_model_the_router_cannot_name() -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + body: Final = {"model": "team-public-name", "messages": [{"role": "user", "content": "hi"}], "tools": _MALFORMED_TOOL_ENTRY} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + assert counter.await_args.kwargs["request"].tools == _MALFORMED_TOOL_ENTRY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dropping", "headers"), + [ + ({"drop_params": True}, {}), + ({"additional_drop_params": ["tools"]}, {}), + ({}, {"user-agent": "claude-cli/2.1.0 (external, cli)"}), + ], + ids=["drop_params", "additional_drop_params", "claude_code_user_agent"], +) +async def test_count_tokens_drops_malformed_tools_when_drop_params_is_set( + dropping: dict[str, object], headers: dict[str, str] +) -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + request: Final = MagicMock() + request.headers = headers + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + messages: Final = [{"role": "user", "content": "hi"}] + body: Final = {"model": "claude-sonnet", "messages": messages, "tools": 5, **dropping} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + counted: Final = counter.await_args.kwargs["request"] + assert counted.tools is None + assert counted.messages == messages + + +@pytest.mark.asyncio +async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_even_with_fallbacks() -> None: + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.proxy._types import UserAPIKeyAuth + + router: Final = Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "sk-test", + "additional_drop_params": ["tools"], + }, + }, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-haiku-5-5", "api_key": "sk-test"}}, + ], + fallbacks=[{"claude-sonnet": ["backup"]}], + ) + request: Final = MagicMock() + request.headers = {} + counter: Final = AsyncMock(return_value=_SEVEN_TOKENS) + body: Final = {"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}], "tools": 5} + + with ( + patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives + patch.object(proxy_server, "llm_router", new=router), # test-quality-ok: module global imported at call time; no injection seam on this route + ): + response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert response == {"input_tokens": 7} + assert counter.await_args.kwargs["request"].tools is None + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index c224cd08663..7de249bfb6e 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -35,7 +35,7 @@ from litellm.proxy.auth.auth_utils import ( is_request_body_safe, ) from litellm.router import Router -from litellm.types.utils import oauth_token_exchange_litellm_params +from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS @@ -163,13 +163,45 @@ def test_every_server_owned_identity_param_is_refused_from_a_request_body(param: ) +def _client_router_config(placement: str, litellm_params: dict[str, str]) -> dict[str, object]: + if placement == "model_list": + return {"model_list": [{"model_name": "gpt-5.4-nano", "litellm_params": litellm_params}]} + return {"default_litellm_params": litellm_params} + + +@pytest.mark.parametrize("param", server_owned_wif_litellm_params) +@pytest.mark.parametrize("placement", ["model_list", "default_litellm_params"]) +def test_every_server_owned_auth_param_is_refused_from_a_client_router_config(param: str, placement: str): + user_config: Final = _client_router_config(placement, {"model": "openai/gpt-5.4-nano", param: "attacker-chosen"}) + with pytest.raises( + ValueError, + match="server-owned workload identity federation or OAuth token exchange parameter", + ): + is_request_body_safe( + request_body={"model": "gpt-5.4-nano", "user_config": user_config}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-5.4-nano", + ) + + @pytest.mark.parametrize( "body", [ {"model": "claude-sonnet-5", "litellm_credential_name": "admin-wif"}, {"model": "claude-sonnet-5", "litellm_params": {"litellm_credential_name": "admin-wif"}}, + { + "model": "claude-sonnet-5", + "user_config": _client_router_config( + "model_list", {"model": "anthropic/claude-sonnet-5", "litellm_credential_name": "admin-wif"} + ), + }, + { + "model": "claude-sonnet-5", + "user_config": _client_router_config("default_litellm_params", {"litellm_credential_name": "admin-wif"}), + }, ], - ids=["top_level", "nested_litellm_params"], + ids=["top_level", "nested_litellm_params", "client_router_model_list", "client_router_defaults"], ) def test_a_request_body_cannot_pick_a_federated_identity_by_credential_name(monkeypatch, body: dict): """Naming a federated credential moves the token exchange onto that credential's federation rule diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index 9f9a6687210..660cc93874d 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -932,6 +932,26 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): assert fake_model not in litellm.models_by_provider["vertex_ai"] +def test_wildcard_lists_a_model_once_when_registered_both_bare_and_provider_prefixed(): + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + bare: Final = "gpt-dedupe-regression" + entry: Final = {"litellm_provider": "openai", "mode": "chat"} + try: + litellm.register_model({bare: entry, f"openai/{bare}": entry}, persist_across_reloads=False) + + listed: Final = get_known_models_from_wildcard("openai/*") + + assert listed.count(f"openai/{bare}") == 1 + finally: + litellm.model_cost.pop(bare, None) + litellm.model_cost.pop(f"openai/{bare}", None) + litellm.open_ai_chat_completion_models.discard(bare) + litellm.open_ai_chat_completion_models.discard(f"openai/{bare}") + litellm.add_known_models(model_cost_map={}) + + def test_get_complete_model_list_drops_no_default_models_sentinel(): from litellm.proxy.auth.model_checks import get_complete_model_list diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index aab2074e49f..15a254dd7a1 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -35,6 +35,7 @@ import logging from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from pathlib import Path from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -521,6 +522,29 @@ async def test_create__fallback_explicit_provider_bypasses_not_found_gate(harnes assert harness.acreate_kwargs()["custom_llm_provider"] == "anthropic" +@pytest.mark.asyncio +async def test_create__fallback_workload_identity_alone_forwards( + harness: Harness, no_openai_creds: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + token_file = tmp_path / "subject_token.jwt" + token_file.write_text("subject-token-from-file") + monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_test123") + monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-test456") + monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file)) + set_body( + harness, + { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + await call_create(harness) + + assert harness.acreate_kwargs()["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_create__fallback_env_key_alone_forwards(harness, monkeypatch): monkeypatch.setattr(litellm, "api_key", None) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py index 43f1a820885..9a4b02ba3ba 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -1,12 +1,81 @@ import asyncio import json +from typing import Final import pytest from fastapi.testclient import TestClient +from prometheus_client import REGISTRY +import litellm from litellm.constants import MAX_SIZE_IN_MEMORY_QUEUE from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.types.services import ServiceTypes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("aggregation_threshold", (2, 3)) +async def test_spend_queue_gauge_tracks_enqueue_aggregation_and_drain( + monkeypatch: pytest.MonkeyPatch, aggregation_threshold: int +) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + queue: Final = SpendUpdateQueue() + queue.MAX_SIZE_IN_MEMORY_QUEUE = aggregation_threshold + response_cost: Final = 0.5 + update: Final = SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, entity_id="queue-user", response_cost=response_cost + ) + + await queue.add_update(update) + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 + + await queue.add_update(update) + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == (1 if aggregation_threshold == 2 else 2) + + await queue.aggregate_queue_updates() + await asyncio.sleep(0) + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 + + transactions: Final = await queue.flush_and_get_aggregated_db_spend_update_transactions() + await asyncio.sleep(0) + assert transactions["user_list_transactions"] == {"queue-user": 2 * response_cost} + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 0 + + +@pytest.mark.asyncio +async def test_spend_queue_gauge_tracks_partial_flush(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.db.db_transaction_queue.base_update_queue as base_queue_module + + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + monkeypatch.setattr(base_queue_module, "MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1) + queue: Final = SpendUpdateQueue() + update: Final = SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, entity_id="queue-user", response_cost=0.5 + ) + await queue.add_update(update) + await queue.add_update(update) + + transactions: Final = await queue.flush_and_get_aggregated_db_spend_update_transactions() + await asyncio.sleep(0) + assert transactions["user_list_transactions"] == {"queue-user": update["response_cost"]} + assert REGISTRY.get_sample_value( + "litellm_in_memory_spend_update_queue_size", + {"in_memory_spend_update_queue": str(ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE)}, + ) == queue.update_queue.qsize() == 1 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 2c91569de69..930309d62ed 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -5,6 +5,7 @@ Test OpenAI Moderation Guardrail import json import os +from pathlib import Path from typing import Final @@ -12,13 +13,22 @@ from unittest.mock import MagicMock, patch import httpx import pytest +import respx +import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers +from litellm.proxy.guardrails.guardrail_hooks.openai import ( + initialize_guardrail as initialize_openai_moderation_guardrail, +) from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( OpenAIModerationGuardrail, ) +from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER, InMemoryGuardrailHandler +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.guardrails import Guardrail, LitellmParams, SupportedGuardrailIntegrations from litellm.types.llms.openai import OpenAIModerationResponse, OpenAIModerationResult @@ -1089,3 +1099,153 @@ async def test_openai_moderation_responses_api_input_field(): input_type="request", ) assert "Violated OpenAI moderation policy" in str(chat_error.value) + + +class _ModerationsEndpointThatRecordsTheAuthorization: + authorization: str | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.authorization = request.headers.get("Authorization") + return httpx.Response( + 200, + json={ + "id": "modr-federated", + "model": "omni-moderation-latest", + "results": [ + {"flagged": False, "categories": {}, "category_scores": {}, "category_applied_input_types": {}} + ], + }, + ) + + +@pytest.mark.asyncio +@respx.mock +async def test_openai_moderation_sends_the_federated_bearer_when_only_workload_identity_is_set( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + token_file: Final = tmp_path / "subject_token.jwt" + token_file.write_text("subject-token") + for name in ("OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.delenv(name, raising=False) + for attribute in ("api_key", "openai_key", "api_base"): + monkeypatch.setattr(litellm, attribute, None) + monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_guardrail") + monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-guardrail") + monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file)) + _workload_identity_auth.cache_clear() + respx.post("https://auth.openai.com/oauth/token").mock( + return_value=httpx.Response(200, json={"access_token": "guardrail-bearer", "expires_in": 3600}) + ) + endpoint: Final = _ModerationsEndpointThatRecordsTheAuthorization() + http_client: Final = AsyncHTTPHandler() + http_client.client = httpx.AsyncClient(transport=httpx.MockTransport(endpoint)) + + guardrail: Final = OpenAIModerationGuardrail(guardrail_name="openai-mod") + guardrail.async_handler = http_client + response: Final = await guardrail.async_make_request("hello") + + assert response.id == "modr-federated" + assert endpoint.authorization == "Bearer guardrail-bearer" + + +@pytest.mark.asyncio +@respx.mock +async def test_openai_moderation_federates_with_the_identity_in_its_own_guardrail_config( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + token_file: Final = tmp_path / "subject_token.jwt" + token_file.write_text("subject-token") + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_IDENTITY_PROVIDER_ID", + "OPENAI_SERVICE_ACCOUNT_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + ): + monkeypatch.delenv(name, raising=False) + for attribute in ("api_key", "openai_key", "api_base"): + monkeypatch.setattr(litellm, attribute, None) + monkeypatch.setenv("GUARDRAIL_IDENTITY_PROVIDER_ID", "idp_guardrail_config") + _workload_identity_auth.cache_clear() + exchange: Final = respx.post("https://auth.openai.com/oauth/token").mock( + return_value=httpx.Response(200, json={"access_token": "guardrail-config-bearer", "expires_in": 3600}) + ) + endpoint: Final = _ModerationsEndpointThatRecordsTheAuthorization() + http_client: Final = AsyncHTTPHandler() + http_client.client = httpx.AsyncClient(transport=httpx.MockTransport(endpoint)) + + guardrail_config: Final = LitellmParams( + guardrail=SupportedGuardrailIntegrations.OPENAI_MODERATION, + mode="pre_call", + openai_identity_provider_id="os.environ/GUARDRAIL_IDENTITY_PROVIDER_ID", + openai_service_account_id="user-guardrail-config", + openai_identity_token_file=str(token_file), + ) + monkeypatch.setattr(litellm, "callbacks", []) + guardrail: Final = initialize_openai_moderation_guardrail( + litellm_params=guardrail_config, + guardrail=Guardrail(guardrail_name="openai-mod", litellm_params=guardrail_config), + ) + guardrail.async_handler = http_client + await guardrail.async_make_request("hello") + + exchange_body: Final = json.loads(exchange.calls.last.request.content) + assert exchange_body["identity_provider_id"] == "idp_guardrail_config" + assert exchange_body["service_account_id"] == "user-guardrail-config" + assert endpoint.authorization == "Bearer guardrail-config-bearer" + + +@pytest.fixture +def no_ambient_openai_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_IDENTITY_PROVIDER_ID", + "OPENAI_SERVICE_ACCOUNT_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + ): + monkeypatch.delenv(name, raising=False) + for attribute in ("api_key", "openai_key", "api_base"): + monkeypatch.setattr(litellm, attribute, None) + monkeypatch.setattr(litellm, "callbacks", []) + + +PARTIAL_IDENTITY_MODERATION: Final = { + "guardrail": SupportedGuardrailIntegrations.OPENAI_MODERATION.value, + "mode": "pre_call", + "openai_identity_provider_id": "idp_without_the_rest", +} + + +@pytest.mark.usefixtures("no_ambient_openai_credentials") +def test_openai_moderation_with_a_partial_identity_is_a_configuration_error_naming_the_gap() -> None: + with pytest.raises(ValueError, match="openai_service_account_id, openai_identity_token_file"): + InMemoryGuardrailHandler().initialize_guardrail( + guardrail={"guardrail_name": "partial-identity-mod", "litellm_params": PARTIAL_IDENTITY_MODERATION}, + source="db", + ) + + +@pytest.mark.usefixtures("no_ambient_openai_credentials") +def test_proxy_boot_skips_an_openai_moderation_with_a_partial_identity() -> None: + IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.clear() + IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.clear() + + init_guardrails_v2( + all_guardrails=[ + {"guardrail_name": "partial-identity-mod", "litellm_params": PARTIAL_IDENTITY_MODERATION}, + { + "guardrail_name": "keyed-mod", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.OPENAI_MODERATION.value, + "mode": "pre_call", + "api_key": "sk-moderation", + }, + }, + ] + ) + + booted: Final = IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.values() + assert {guardrail["guardrail_name"] for guardrail in booted} == {"keyed-mod"} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 19eb3c71a5b..f50da1bff64 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -5,6 +5,8 @@ This test file tests the Generic Guardrail API implementation, specifically focusing on metadata extraction and passing. """ +import json +import logging import os from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -16,6 +18,7 @@ import litellm from litellm import ModelResponse from litellm._version import version as litellm_version from litellm.exceptions import GuardrailRaisedException, Timeout +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPI, @@ -38,6 +41,38 @@ def generic_guardrail(): ) +@pytest.fixture +def raw_rows_guardrail(): + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + headers={"Authorization": "Bearer test-key"}, + guardrail_name="test-generic-guardrail-raw-rows", + event_hook="pre_call", + default_on=True, + forward_messages_as_sent=True, + ) + + +async def _structured_messages_sent(guardrail: GenericGuardrailAPI, rows: list[dict]) -> list: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: records every body the guardrail endpoint received + + async def answer(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.content) + return httpx.Response(200, json={"action": "NONE"}, request=request) + + endpoint: Final = AsyncHTTPHandler(transport=httpx.MockTransport(answer)) + try: + with patch.object(guardrail, "async_handler", endpoint): + await guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": rows}, + request_data={}, + input_type="request", + ) + finally: + await endpoint.close() + return [json.loads(body)["structured_messages"] for body in sent_bodies] + + @pytest.fixture def mock_user_api_key_dict(): """Create a mock UserAPIKeyAuth object""" @@ -638,6 +673,157 @@ class TestStructuredMessagesInResponse: assert returned_rows[1] is tool_call_row assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("row", "openai_row"), + [ + ( + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + {"role": "user", "content": "Look up 123-45-6789 for me."}, + ), + ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "index": 0, + } + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + ), + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAA="}} + ], + }, + {"role": "user", "content": []}, + ), + ], + ids=["extra_key", "tool_call_index", "image_block"], + ) + async def test_rows_reach_the_guardrail_as_openai_chat_rows_by_default( + self, generic_guardrail: GenericGuardrailAPI, row, openai_row + ): + assert await _structured_messages_sent(generic_guardrail, [row]) == [[openai_row]] + + @pytest.mark.asyncio + async def test_a_returned_row_that_differs_from_the_shown_row_is_written_back_by_default( + self, generic_guardrail, mock_request_data_input + ): + sent_row = {"role": "user", "content": "my ssn is 123-45-6789", "name": "pat"} + + def answer_with_the_row_as_the_client_sent_it(url, json, headers, **_kwargs): + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["my ssn is [SSN]"], + "structured_messages": [sent_row], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object( + generic_guardrail.async_handler, "post", side_effect=answer_with_the_row_as_the_client_sent_it + ): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"], "structured_messages": [sent_row]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == [sent_row] + + @pytest.mark.asyncio + async def test_rows_echoed_back_with_one_rewrite_write_back_the_rewrite_only( + self, raw_rows_guardrail, mock_request_data_input + ): + tool_call_row = { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + } + original_rows = [ + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + tool_call_row, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + ] + + def echo_with_tool_output_redacted(url, json, headers, **_kwargs): + shown_rows = json["structured_messages"] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Look up 123-45-6789 for me."], + "structured_messages": [ + shown_rows[0], + shown_rows[1], + {**shown_rows[2], "content": '{"ssn": "[REDACTED]"}'}, + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(raw_rows_guardrail.async_handler, "post", side_effect=echo_with_tool_output_redacted): + guardrailed_inputs = await raw_rows_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + returned_rows = guardrailed_inputs["structured_messages"] + assert returned_rows[0] == original_rows[0] + assert returned_rows[1] == tool_call_row + assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "row", + [ + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "123-45-6789"}]}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "lookup", "input": {}}]}, + { + "role": "user", + "content": [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAA="}}], + }, + { + "role": "user", + "content": [ + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "ssn"}} + ], + }, + ], + ids=["extra_key", "tool_call_index", "tool_result_block", "tool_use_block", "image_block", "document_block"], + ) + async def test_rows_reach_the_guardrail_verbatim_when_forwarded_as_sent( + self, raw_rows_guardrail: GenericGuardrailAPI, row + ): + """Rows are the provider's own shape, so the guardrail must see every key and + content block as sent, not a copy filtered through the chat completion types.""" + assert await _structured_messages_sent(raw_rows_guardrail, [row]) == [[row]] + @pytest.mark.asyncio async def test_rows_all_echoed_back_as_shown_leave_the_rewrite_to_texts( self, generic_guardrail, mock_request_data_input @@ -672,6 +858,126 @@ class TestStructuredMessagesInResponse: assert "structured_messages" not in guardrailed_inputs assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @pytest.mark.asyncio + async def test_rows_echoed_without_the_keys_the_server_does_not_know_leave_the_rewrite_to_texts( + self, raw_rows_guardrail, mock_request_data_input + ): + """A server that parses rows with its own message types echoes them back without + the keys those types lack, and its texts redaction must still reach the LLM.""" + original_rows = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + "refusal": None, + "annotations": [], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + {"role": "user", "content": "my ssn is 123-45-6789"}, + ] + + def echo_known_keys_and_redact_texts(url, json, headers, **_kwargs): + rows = json["structured_messages"] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": [text.replace("123-45-6789", "[SSN]") for text in json["texts"]], + "structured_messages": [ + { + "role": "assistant", + "content": None, + "refusal": None, + "tool_calls": [ + {key: value for key, value in rows[0]["tool_calls"][0].items() if key != "index"} + ], + }, + rows[1], + rows[2], + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(raw_rows_guardrail.async_handler, "post", side_effect=echo_known_keys_and_redact_texts): + guardrailed_inputs = await raw_rows_guardrail.apply_guardrail( + inputs={"texts": ["found", "my ssn is 123-45-6789"], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["found", "my ssn is [SSN]"] + + @pytest.mark.asyncio + async def test_a_row_rewrite_echoes_the_other_rows_back_as_their_original_rows( + self, raw_rows_guardrail, mock_request_data_input + ): + assistant_row = { + "role": "assistant", + "content": "Looking it up.", + "annotations": [{"type": "url_citation", "url_citation": {"url": "https://example.com"}}], + } + original_rows = [assistant_row, {"role": "user", "content": "my ssn is 123-45-6789"}] + redacted_row = {"role": "user", "content": "my ssn is [SSN]"} + + def redact_one_row(url, json, headers, **_kwargs): + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "structured_messages": [{"role": "assistant", "content": "Looking it up."}, redacted_row], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(raw_rows_guardrail.async_handler, "post", side_effect=redact_one_row): + guardrailed_inputs = await raw_rows_guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == [assistant_row, redacted_row] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "returned_row", + [ + {"role": "assistant", "content": None}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + "annotations": [], + "refusal": "blocked", + }, + ], + ids=["drops_a_declared_key", "changes_a_value"], + ) + async def test_a_row_that_loses_a_declared_key_or_changes_a_value_is_a_rewrite( + self, raw_rows_guardrail, mock_request_data_input, returned_row + ): + original_row = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + "annotations": [], + "refusal": None, + } + answer = MagicMock() + answer.json.return_value = {"action": "GUARDRAIL_INTERVENED", "structured_messages": [returned_row]} + answer.raise_for_status = MagicMock() + + with patch.object(raw_rows_guardrail.async_handler, "post", return_value=answer): + guardrailed_inputs = await raw_rows_guardrail.apply_guardrail( + inputs={"texts": [], "structured_messages": [original_row]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == [returned_row] + @pytest.mark.asyncio @pytest.mark.parametrize( "structured_messages", @@ -1253,6 +1559,40 @@ class TestGenericGuardrailAPIStreamingConfig: ) assert guardrail.streaming_transform_mode == "incremental_diff" + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("configured", "forwards_as_sent"), + [ + pytest.param({}, False, id="unset"), + pytest.param({"forward_messages_as_sent": ""}, False, id="blank"), + pytest.param({"forward_messages_as_sent": False}, False, id="false"), + pytest.param({"forward_messages_as_sent": "true"}, True, id="string-true"), + pytest.param({"forward_messages_as_sent": True}, True, id="true"), + ], + ) + async def test_forward_messages_as_sent_config_decides_which_rows_reach_the_guardrail( + self, configured, forwards_as_sent + ): + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams.model_validate( + { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "api_base": "https://api.test.guardrail.com", + **configured, + } + ) + with patch("litellm.logging_callback_manager.add_litellm_callback"): + guardrail = initialize_guardrail(litellm_params, {"guardrail_name": "test-forward-as-sent"}) + + row = {"role": "user", "content": "hi", "name": "pat"} + expected_row = row if forwards_as_sent else {"role": "user", "content": "hi"} + assert await _structured_messages_sent(guardrail, [row]) == [[expected_row]] + def test_initialize_guardrail_forwards_streaming_flags(self): from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( initialize_guardrail, @@ -1945,6 +2285,34 @@ class TestToolSupport: assert forwarded_tools == tools + @pytest.mark.asyncio + async def test_provider_native_tools_without_a_type_reach_the_guardrail_and_come_back_unchanged( + self, generic_guardrail + ): + """Gemini tools are declared as function_declarations, googleSearch and the + like with no type key, so they must be scanned, and a server that echoes + them back must not write a type key into the request to the provider.""" + tools = [ + {"function_declarations": [{"name": "get_weather", "parameters": {"type": "object", "properties": {}}}]}, + {"googleSearch": {}}, + ] + + def echo_tools(url, json, headers, **_kwargs): + answer = MagicMock() + answer.json.return_value = {"action": "NONE", "texts": ["hi"], "tools": json["tools"]} + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_tools) as mock_post: + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["hi"], "tools": tools}, + request_data={}, + input_type="request", + ) + + assert mock_post.call_args.kwargs["json"]["tools"] == tools + assert guardrailed_inputs["tools"] == tools + class TestFailOnError: """Test fail_on_error: complete fail-open on any guardrail error""" @@ -2032,6 +2400,59 @@ class TestFailOnError: assert result == {"texts": ["model output"]} + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "error_type", "pre_pr_text"), + [ + ( + Timeout(message="guardrail took too long", model="gpt-5", llm_provider="openai"), + "Timeout", + "Proceeding without guardrail. guardrail_name=", + ), + ( + httpx.ConnectError("connection refused", request=MagicMock()), + "ConnectError", + "Proceeding without guardrail. guardrail_name=", + ), + ( + MaskedHTTPStatusError( + httpx.HTTPStatusError( + "bad gateway", + request=httpx.Request("POST", "http://guardrail.invalid/scan"), + response=httpx.Response(502, request=httpx.Request("POST", "http://guardrail.invalid/scan")), + ) + ), + "HTTPStatusError", + "Proceeding without guardrail. http_status_code=502 guardrail_name=", + ), + ], + ids=["timeout", "connect_error", "http_status"], + ) + async def test_fail_open_log_line_names_the_error_type( + self, fail_open_guardrail, caplog: pytest.LogCaptureFixture, error, error_type, pre_pr_text + ): + """An operator grepping the fail-open line must be able to tell a timeout + from a refused connection from an upstream status without the traceback.""" + from litellm._logging import verbose_proxy_logger + + verbose_proxy_logger.addHandler(caplog.handler) + try: + with patch.object(fail_open_guardrail.async_handler, "post", side_effect=error): + await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + fail_open_lines: Final = [ + record.getMessage() for record in caplog.records if record.levelno == logging.CRITICAL + ] + assert fail_open_lines + assert all(f"error_type={error_type}" in line.split() for line in fail_open_lines) + assert all(pre_pr_text in line for line in fail_open_lines) + @pytest.mark.asyncio async def test_response_path_valid_block_still_blocks(self, fail_open_guardrail): """On the response path too, a valid BLOCKED decision raises despite fail_on_error=False.""" diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 908a49946bb..2e968ec04f5 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -953,6 +953,17 @@ async def test_provider_specific_params_includes_embedding_toggle(): assert field["default_value"] is False +@pytest.mark.asyncio +async def test_provider_specific_params_offers_forward_messages_as_sent_off_by_default(): + from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params + + provider_params = await get_provider_specific_params() + + field = provider_params["generic_guardrail_api"]["optional_params"]["fields"]["forward_messages_as_sent"] + assert field["type"] == "boolean" + assert field["default_value"] is False + + @pytest.mark.asyncio async def test_provider_specific_params_exposes_bedrock_streaming_flags(): from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params diff --git a/tests/unit/proxy/hooks/test_prompt_cache_observer.py b/tests/unit/proxy/hooks/test_prompt_cache_observer.py index 3354aa457b7..ea89bb8d4a2 100644 --- a/tests/unit/proxy/hooks/test_prompt_cache_observer.py +++ b/tests/unit/proxy/hooks/test_prompt_cache_observer.py @@ -201,6 +201,10 @@ class RecordingObserver(PromptCacheObserver): await super().async_log_success_event(kwargs, response_obj, start_time, end_time) self.finished.set() + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + self.finished.set() + def native_response(): return { diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index beffe454b8d..ad75344afba 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3758,6 +3758,43 @@ class TestUpdateDBModelClearPricing: assert "cache_creation_input_token_cost" not in params assert "cache_creation_input_token_cost" not in info + def test_clear_cache_write_1h_cost_removes_from_both_blobs(self): + """The dashboard clears the 1 hour cache-write tier with an explicit null, which only + takes effect for fields in SPECIAL_MODEL_INFO_PARAMS.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_db_model, + ) + from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelInfo, + updateLiteLLMParams, + ) + + db_model = Deployment( + model_name="anthropic/*", + litellm_params=LiteLLM_Params( + model="anthropic/*", + cache_creation_input_token_cost=0.00000375, + cache_creation_input_token_cost_above_1hr=0.000006, + ), + model_info=ModelInfo(id="dep-cache-write-1h-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(cache_creation_input_token_cost_above_1hr=None) + ), + ) + + params = json.loads(result["litellm_params"]) + info = json.loads(result["model_info"]) + assert "cache_creation_input_token_cost_above_1hr" not in params + assert "cache_creation_input_token_cost_above_1hr" not in info + assert params["cache_creation_input_token_cost"] == 0.00000375 + assert info["cache_creation_input_token_cost"] == 0.00000375 + def test_clear_cache_read_preserves_other_pricing(self): """Clearing cache_read must not touch input/output cost overrides.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 765e0506990..873182a5145 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -2,7 +2,7 @@ import asyncio import json import logging import os -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace from typing import Final @@ -17,7 +17,7 @@ from fastapi import HTTPException, Request import litellm from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse +from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -26,9 +26,11 @@ from litellm.proxy.management_endpoints.ui_sso import ( MicrosoftSSOHandler, SSOAuthenticationHandler, _complete_cli_sso_callback_session, + _insert_sso_user_or_find_concurrent_insert, _setup_team_mappings, _sync_user_role_from_jwt_role_map, cli_sso_callback, + insert_sso_user, normalize_email, process_sso_jwt_access_token, sso_lookup_user_ids, @@ -1425,6 +1427,82 @@ async def test_upsert_sso_user_no_role_in_sso_response(): assert "user_role" not in call_args.kwargs["data"] +class _UserTable: + def __init__(self, *rows: LiteLLM_UserTable) -> None: + self.rows: dict[str, LiteLLM_UserTable] = {row.user_id: row for row in rows} + + async def find_first(self, where: dict[str, dict[str, str]]) -> LiteLLM_UserTable | None: + ((field, condition),) = where.items() + + def matches(stored: object) -> bool: + if condition.get("mode") == "insensitive": + return str(stored).lower() == condition["equals"].lower() + return stored == condition["equals"] + + return next((row for row in self.rows.values() if matches(getattr(row, field))), None) + + async def find_unique(self, where: dict[str, str]) -> LiteLLM_UserTable | None: + return self.rows.get(where["user_id"]) + + +def _users_db(table: _UserTable) -> SimpleNamespace: + database: Final = SimpleNamespace(litellm_usertable=table) + return SimpleNamespace(db=database, writer_db=database) + + +async def _first_sign_in(table: _UserTable, **join: object) -> NewUserResponse | LiteLLM_UserTable: + prisma: Final = _users_db(table) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + return await _insert_sso_user_or_find_concurrent_insert( + insert=insert_sso_user( + result_openid=CustomOpenID(id="sso-subject", email="person@example.com", team_ids=[]), + user_defined_values={ + "user_id": "sso-subject", + "user_email": "person@example.com", + "user_role": None, + "max_budget": None, + "budget_duration": None, + "models": [], + }, + ), + user_id="sso-subject", + prisma_client=prisma, + **join, + ) + + +@pytest.mark.asyncio +async def test_first_sign_in_signs_in_to_the_account_a_concurrent_first_sign_in_created(): + winner: Final = LiteLLM_UserTable(user_id="sso-subject", user_email="person@example.com", user_role="internal_user") + + signed_in: Final = await _first_sign_in(_UserTable(winner)) + + assert signed_in == winner + + +@pytest.mark.asyncio +async def test_first_sign_in_does_not_sign_in_to_another_account_holding_the_email(): + other: Final = LiteLLM_UserTable(user_id="someone-else", user_email="Person@example.com", user_role="proxy_admin") + + with pytest.raises(ProxyException, match=r"User with email Person@example\.com already exists"): + await _first_sign_in(_UserTable(other)) + + +@pytest.mark.asyncio +async def test_first_sign_in_joins_default_teams_the_concurrent_first_sign_in_has_not_joined_yet(monkeypatch): + monkeypatch.setattr(litellm, "default_internal_user_params", {"teams": ["team-joined", "team-pending"]}) + table: Final = _UserTable( + LiteLLM_UserTable(user_id="sso-subject", user_email="person@example.com", teams=["team-joined"]) + ) + + async def join(user: LiteLLM_UserTable, teams: Sequence[object], prisma_client: object) -> None: + table.rows[user.user_id] = user.model_copy(update={"teams": [*user.teams, *map(str, teams)]}) + + signed_in: Final = await _first_sign_in(table, join_default_teams=join) + + assert signed_in.teams == ["team-joined", "team-pending"] + + def test_get_user_email_and_id_extracts_microsoft_role(): """ Test that _get_user_email_and_id_from_result extracts user_role from Microsoft SSO. diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 88c5e2d2c1b..4770ecd5485 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -5,11 +5,10 @@ from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx -import litellm import pytest from pydantic import ValidationError - +import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data @@ -2644,6 +2643,58 @@ class TestRecordPartialUsageForFailure: assert "response_cost" not in logging_obj.model_call_details +class _FailureRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failure_kwargs = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.failure_kwargs.append(kwargs) + + +class TestLogErrorFrameAsFailure: + """A stream closed at the provider's `event: error` frame is one failed attempt: its failure log bills the + usage streamed before the frame on the attempt-scoped copy its caller hands it, and the request's shared + logging state stays as it was, so the attempt the router opens next logs with nothing left behind.""" + + @pytest.mark.asyncio + async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self): + recorder = _FailureRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-error-frame-failure", + function_id="test-error-frame-failure", + dynamic_async_failure_callbacks=[recorder], + ) + chunks = [ + *TestRecordPartialUsageForFailure()._interrupted_chunks()[:1], + TestRecordPartialUsageForFailure._sse( + "error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + ), + ] + + await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( + attempt_logging_obj=logging_obj.attempt_scoped_copy(), + request_body={"model": "claude-sonnet-5", "stream": True}, + all_chunks=chunks, + error_event=("overloaded_error", "Overloaded", 529), + ) + + assert len(recorder.failure_kwargs) == 1 + failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] + assert failure_payload["status"] == "failure" + assert failure_payload["prompt_tokens"] == 52 + assert failure_payload["response_cost"] > 0 + assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError) + assert "combined_usage_object" not in logging_obj.model_call_details + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.should_run_logging("async_failure") + + @pytest.mark.parametrize( ("all_chunks", "interrupted"), [ diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index b278d7dea01..8ce41129aec 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -18,6 +18,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthr from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.anthropic.pass_through.messages.streaming_iterator import aclose_superseded_attempt from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) @@ -749,6 +750,294 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: + return _anthropic_sse( + "message_start", + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": input_tokens, "output_tokens": 1}, + }, + }, + ) + + +def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: + return ( + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "pong"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}}, + ), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ) + + +async def _settle(recorded: list[dict[str, object]]) -> None: + for _ in range(300): + if recorded: + return + await asyncio.sleep(0.01) + + +def _anthropic_overloaded_error_frame() -> bytes: + return _anthropic_sse("error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}) + + +def _anthropic_stream_of(frames: tuple[bytes, ...]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for frame in frames: + yield frame + + mock.aiter_bytes = _aiter_bytes + return mock + + +@pytest.mark.parametrize( + "frames, billed", + [ + pytest.param( + (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start" + ), + pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"), + ], +) +@pytest.mark.asyncio +async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_closed_at_the_frame(frames, billed): + """The router retries a /v1/messages stream whose provider sent `event: error` before any content by closing + it at the frame. That attempt never completed: it logs once as a failure carrying whatever usage the provider + had reported, and the request's one-shot success log stays free for the retry that shares the logging object.""" + recorder = _EventRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-error-frame", + function_id="test-error-frame", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + stream = PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + received = [] + async for chunk in stream: + received.append(chunk) + if b"event: error" in chunk: + break + await aclose_superseded_attempt(stream) + await _settle(recorder.failure_kwargs) + + assert received == list(frames) + assert recorder.success_kwargs == [] + assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True + assert "combined_usage_object" not in logging_obj.model_call_details + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.should_run_logging("async_failure") + assert len(recorder.failure_kwargs) == 1 + failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] + assert failure_payload["status"] == "failure" + assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError) + if billed: + assert failure_payload["prompt_tokens"] == 52 + assert failure_payload["response_cost"] > 0 + else: + assert failure_payload["response_cost"] == 0 + + +@pytest.mark.asyncio +async def test_chunk_processor_logs_a_stream_that_ended_at_the_provider_error_frame_through_the_success_route(): + """A provider error frame that ends a stream nobody closed early is the request's final answer, and that + stream keeps logging the way it always has: once, through the success route, with the usage it carried.""" + recorder = _EventRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-final-error-frame", + function_id="test-final-error-frame", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + frames = (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()) + received = [ + chunk + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + ] + await _settle(recorder.success_kwargs) + + assert received == list(frames) + assert recorder.failure_kwargs == [] + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert success_payload["prompt_tokens"] == 52 + assert logging_obj.model_call_details["prompt_cache_response_complete"] is False + + +@pytest.mark.asyncio +async def test_chunk_processor_bills_a_stream_the_client_closed_at_the_provider_error_frame(): + """The Anthropic SDK raises on a provider `event: error` frame and drops the connection, so a non-retried + stream is closed at the frame by the client, not by a retry taking over. Its usage must still reach the + success route, or the tokens the provider billed never reach the key's spend.""" + recorder = _EventRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-client-closed-at-error-frame", + function_id="test-client-closed-at-error-frame", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + stream = PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of((_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())), + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + async for chunk in stream: + if b"event: error" in chunk: + break + await stream.aclose() + await _settle(recorder.success_kwargs) + + assert recorder.failure_kwargs == [] + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert success_payload["prompt_tokens"] == 52 + + +@pytest.mark.asyncio +async def test_superseded_attempt_failure_log_keeps_its_own_deployment_when_the_retry_rebinds_the_logging_object(): + """The router rebinds the shared logging object to the retry's deployment right after closing the superseded + stream, before the logging worker runs. The failure log of the superseded attempt must still name the + deployment that served it, or the router's cooldown lands on the deployment serving the retry.""" + recorder = _EventRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-superseded-deployment", + function_id="test-superseded-deployment", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="claude-sonnet-5", litellm_params={"model_info": {"id": "deployment-1"}}, optional_params={} + ) + superseded = PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of((_anthropic_message_start_frame(), _anthropic_overloaded_error_frame())), + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + async for chunk in superseded: + if b"event: error" in chunk: + break + await aclose_superseded_attempt(superseded) + logging_obj.update_environment_variables( + model="claude-sonnet-5-retry", litellm_params={"model_info": {"id": "deployment-2"}}, optional_params={} + ) + await _settle(recorder.failure_kwargs) + + failure_kwargs = recorder.failure_kwargs[0] + assert failure_kwargs["litellm_params"]["model_info"]["id"] == "deployment-1" + assert failure_kwargs["model"] == "claude-sonnet-5" + assert logging_obj.model_call_details["litellm_params"]["model_info"]["id"] == "deployment-2" + + +@pytest.mark.asyncio +async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame(): + """Both attempts of a retried /v1/messages stream log through the request's one logging object. The + superseded attempt's failure log must not leave the usage it consumed behind, or the retry's success + log bills the failed attempt's tokens instead of its own.""" + recorder = _EventRecorder() + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="test-retry-own-usage", + function_id="test-retry-own-usage", + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + def attempt(frames: tuple[bytes, ...]): + return PassThroughStreamingHandler.chunk_processor( + response=_anthropic_stream_of(frames), + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + superseded = attempt((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + async for chunk in superseded: + if b"event: error" in chunk: + break + await aclose_superseded_attempt(superseded) + await _settle(recorder.failure_kwargs) + retry = attempt((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] + assert failure_payload["prompt_tokens"] == 52 + success_payload = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 3fa930f9855..8a01feb7492 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3459,10 +3459,11 @@ class TestResponsesInputTokens: body: dict[str, Any], path: str = "/v1/responses/input_tokens", counter: AsyncMock | None = None, + llm_router: litellm.Router | None = None, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.response_api_endpoints.endpoints import _proxy_token_counter + from litellm.proxy.response_api_endpoints.endpoints import _proxy_llm_router, _proxy_token_counter from litellm.types.utils import TokenCountResponse token_counter_mock = ( @@ -3479,6 +3480,7 @@ class TestResponsesInputTokens: ) app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-test", request_route=path) app.dependency_overrides[_proxy_token_counter] = lambda: token_counter_mock + app.dependency_overrides[_proxy_llm_router] = lambda: llm_router try: client = TestClient(app) response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876"}) @@ -3486,6 +3488,7 @@ class TestResponsesInputTokens: finally: app.dependency_overrides.pop(user_api_key_auth, None) app.dependency_overrides.pop(_proxy_token_counter, None) + app.dependency_overrides.pop(_proxy_llm_router, None) def test_string_input_returns_openai_input_tokens_shape(self): response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"}) @@ -3599,6 +3602,48 @@ class TestResponsesInputTokens: assert response.status_code == 400, response.text error = response.json()["error"] assert error["type"] == "invalid_request_error" + assert error["param"] == "tools" + assert error["message"].startswith("tools: ") + counter.assert_not_awaited() + + def test_invalid_tools_are_dropped_under_drop_params(self) -> None: + response, counter = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi", "tools": [5, {"type": "web_search"}], "drop_params": True} + ) + + assert response.status_code == 200, response.text + assert counter.await_args.kwargs["request"].tools == [{"type": "web_search"}] + + def test_malformed_tools_follow_the_counted_deployment_even_when_fallbacks_are_configured(self) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.2", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]}, + }, + {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}}, + ], + fallbacks=[{"gpt-5.2": ["backup"]}], + ) + + response, counter = self._post_input_tokens( + {"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router + ) + + assert response.status_code == 200, response.text + assert counter.await_args.kwargs["request"].tools is None + + def test_malformed_tools_for_a_model_outside_the_router_get_a_clean_400(self) -> None: + router: Final = litellm.Router( + model_list=[{"model_name": "gpt-5.2", "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test"}}] + ) + + response, counter = self._post_input_tokens( + {"model": "openai/gpt-5.4-mini", "input": "hi", "tools": 5}, llm_router=router + ) + + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "tools" counter.assert_not_awaited() def test_provider_error_maps_status_code(self): diff --git a/tests/unit/proxy/test_tools_allowlist_enforcement.py b/tests/unit/proxy/test_tools_allowlist_enforcement.py index e855ef47dd5..e306f7be4f8 100644 --- a/tests/unit/proxy/test_tools_allowlist_enforcement.py +++ b/tests/unit/proxy/test_tools_allowlist_enforcement.py @@ -145,6 +145,21 @@ class TestExtractRequestToolNames: data = {"tools": [{"type": "function", "function": {"name": "x"}}]} assert extract_request_tool_names("/v1/embeddings", data) == [] + @pytest.mark.parametrize( + ("route", "data"), + [ + ("/v1/chat/completions", {"tools": 5}), + ("/v1/chat/completions", {"functions": 5}), + ("/v1/responses", {"tools": 5}), + ("/v1/messages", {"tools": 5}), + ("/generate_content", {"tools": 5}), + ("/generate_content", {"tools": [{"functionDeclarations": 5}]}), + ], + ids=["chat_tools", "chat_functions", "responses", "messages", "generate_content", "function_declarations"], + ) + def test_malformed_tool_lists_yield_no_names(self, route: str, data: dict[str, object]) -> None: + assert extract_request_tool_names(route, data) == [] + class TestCheckToolsAllowlist: """Test allowlist enforcement in auth (no DB in hot path).""" diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 84aa1798da4..e83efa21dcc 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio +from collections.abc import Callable from datetime import datetime from types import MappingProxyType from typing import Final @@ -921,6 +922,48 @@ async def test_post_call_failure_hook_falls_back_to_body_shape_without_a_route(p assert request_data["call_type"] == "acompletion" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, body", + [ + ("anthropic_messages", {"model": "m", "messages": [{"role": "user", "content": "hi"}]}), + ("aresponses", {"model": "m", "input": [{"role": "user", "content": "hi"}]}), + ("acompact_responses", {"model": "m", "input": [{"role": "user", "content": "hi"}]}), + ], +) +async def test_post_call_failure_hook_keeps_the_call_type_request_setup_logged( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + call_type: str, + body: dict[str, object], +) -> None: + """A request rejected after request setup already carries its endpoint's call type, so the body + shape must not relabel a /v1/messages rejection as a chat completion or a list-input /v1/responses + rejection as an embedding.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + proxy_logging.alert_types = [] + request_data: Final = { + **body, + "litellm_logging_obj": LiteLLMLoggingObj( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=call_type, + start_time=datetime.now(), + litellm_call_id=f"test_keeps_{call_type}", + function_id=f"test_keeps_{call_type}", + ), + "metadata": {}, + } + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="tools must be a list"), + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + ) + assert request_data["call_type"] == call_type + + @pytest.mark.asyncio @pytest.mark.parametrize("route", ["/v1/files", "/files/file-abc", "/v1/containers"]) async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes( diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 41179967a7f..527d3d599b1 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -2,17 +2,20 @@ import asyncio import json import time from collections.abc import Awaitable, Callable, Mapping +from pathlib import Path from types import TracebackType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.models.credentials import CredentialItem from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model @@ -1052,3 +1055,124 @@ async def test_voice_transcript_blocked_by_guardrail(): finally: litellm.callbacks = [] + + +@pytest.fixture +def federated_deployment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> dict[str, str]: + token_file: Final = tmp_path / "subject_token.jwt" + token_file.write_text("subject-token") + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_IDENTITY_PROVIDER_ID", + "OPENAI_SERVICE_ACCOUNT_ID", + "OPENAI_IDENTITY_TOKEN_FILE", + ): + monkeypatch.delenv(name, raising=False) + for attribute in ("api_key", "openai_key", "api_base"): + monkeypatch.setattr(litellm, attribute, None) + _workload_identity_auth.cache_clear() + return { + "openai_identity_provider_id": "idp_realtime", + "openai_service_account_id": "user-realtime", + "openai_identity_token_file": str(token_file), + } + + +def _mock_token_exchange() -> None: + respx.post("https://auth.openai.com/oauth/token").mock( + return_value=httpx.Response(200, json={"access_token": "federated-bearer", "expires_in": 3600}) + ) + + +class _ConnectThatStopsAfterCapturingTheAuthorization: + authorization: str | None = None + + def __call__( + self, url: str, *, additional_headers: Mapping[str, str], **kwargs: object + ) -> "_ConnectThatStopsAfterCapturingTheAuthorization": + self.authorization = additional_headers.get("Authorization") + return self + + async def __aenter__(self) -> None: + raise RuntimeError("backend headers captured, nothing to bridge") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "live_session_start", + [None, {"type": "session.start", "session": {"model": "gpt-live-1"}}], + ids=["realtime", "live_session"], +) +@respx.mock +async def test_openai_websocket_sessions_send_the_federated_bearer_for_a_keyless_deployment( + federated_deployment: dict[str, str], live_session_start: Mapping[str, object] | None +) -> None: + _mock_token_exchange() + connect: Final = _ConnectThatStopsAfterCapturingTheAuthorization() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="gpt-realtime", + custom_llm_provider="openai", + websocket=MagicMock(close=AsyncMock()), + litellm_logging_obj=FakeLogging(), + live_session_start=live_session_start, + **federated_deployment, + ) + assert connect.authorization == "Bearer federated-bearer" + + +class _GatewayThatRecordsTheAuthorization: + authorization: str | None = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.authorization = request.headers.get("Authorization") + return httpx.Response(200, json={"value": "ek_test"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [realtime_main.acreate_realtime_client_secret, realtime_main.acreate_realtime_transcription_session], + ids=["client_secrets", "transcription_sessions"], +) +@respx.mock +async def test_openai_realtime_http_routes_send_the_federated_bearer_for_a_keyless_deployment( + federated_deployment: dict[str, str], route: Callable[..., Awaitable[object]] +) -> None: + _mock_token_exchange() + gateway: Final = _GatewayThatRecordsTheAuthorization() + await route.__wrapped__( + model="gpt-realtime", + custom_llm_provider="openai", + litellm_logging_obj=FakeLogging(), + client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)), + **federated_deployment, + ) + assert gateway.authorization == "Bearer federated-bearer" + + +@pytest.mark.asyncio +@respx.mock +async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_keyless_deployment( + federated_deployment: dict[str, str], +) -> None: + _mock_token_exchange() + connect: Final = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main.realtime_health_check( + model="gpt-realtime", + custom_llm_provider="openai", + api_key=None, + model_params={"model": "openai/gpt-realtime", **federated_deployment}, + ) + assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer federated-bearer"} diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 253850bbd79..e11c5d6c2f3 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -5357,6 +5357,7 @@ class TestRouterPreRoutingSharedAliasName: "api_key": "sk-marker", "api_base": "https://marker.example/v1", "api_version": "2024-01-01", + "responses_api_version": "v1", }, } router = Router(model_list=[marker_with_connection_params, self._tier_entry()]) @@ -5372,6 +5373,7 @@ class TestRouterPreRoutingSharedAliasName: assert "api_base" not in request_kwargs assert "api_key" not in request_kwargs assert "api_version" not in request_kwargs + assert "responses_api_version" not in request_kwargs assert request_kwargs["drop_params"] is True @pytest.mark.asyncio diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index af8fca72252..3532870df1c 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -287,6 +287,19 @@ class TestFilterWebSearchDeployments: assert "deployment-2" in result_ids assert "deployment-3" not in result_ids + def test_web_search_in_a_tuple_filters_unsupported(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = filter_web_search_deployments(sample_deployments, {"tools": ({"type": "web_search"},)}) + assert [d["model_info"]["id"] for d in result] == ["deployment-1", "deployment-2"] + + @pytest.mark.parametrize("tool_type", [{}, [], None, 5], ids=["object", "list", "null", "int"]) + def test_a_non_string_tool_type_is_no_web_search_request( + self, sample_deployments: list[dict[str, object]], tool_type: object + ) -> None: + result: Final = filter_web_search_deployments( + sample_deployments, {"tools": [{"type": tool_type}, {"type": "function"}]} + ) + assert result == sample_deployments + def test_web_search_preview_filters_unsupported(self, sample_deployments): """web_search_preview type should also trigger filtering""" request_kwargs = {"tools": [{"type": "web_search_preview"}]} diff --git a/tests/unit/test_assert_ci_coverage.py b/tests/unit/test_assert_ci_coverage.py index c0e5339c925..87c56f1a42b 100644 --- a/tests/unit/test_assert_ci_coverage.py +++ b/tests/unit/test_assert_ci_coverage.py @@ -9,6 +9,8 @@ the question neither covers: whether the job that globs a file then deselects it """ import importlib.util +import shutil +import subprocess import sys from pathlib import Path from typing import Final @@ -71,17 +73,27 @@ def test_integration_groups_require_exclusive_scheduled_circleci_owner(tmp_path: workflow.write_text(yaml.safe_dump({"jobs": {"tests": {"steps": [{"run": "pytest tests/integration"}]}}})) _, findings = coverage._integration_ownership(tmp_path) assert [(finding.subject, finding.detail) for finding in findings] == [ - (test_path, "integration contract is also selected by GitHub Actions") + (test_path, "integration file outside GITHUB_FILES is selected by GitHub Actions") ] - github_path: Final = "tests/integration/management/test_github_contract.py" + github_path: Final = "tests/integration/database/image/test_github_contract.py" + (tmp_path / github_path).parent.mkdir(parents=True) (tmp_path / github_path).write_text("def test_contract(): pass\n") runner: Final = tmp_path / "tests/integration/run.py" runner.write_text(runner.read_text() + f"GITHUB_FILES: Final = frozenset({{{github_path!r}}})\n") - workflow.write_text(yaml.safe_dump({"jobs": {"tests": {"steps": [{"run": f"pytest {github_path}"}]}}})) + workflow.unlink() + dedicated: Final = tmp_path / ".github/workflows/image-scan.yml" + dedicated.write_text(yaml.safe_dump({"jobs": {"tests": {"steps": [{"run": f"pytest {github_path}"}]}}})) github_owned, github_findings = coverage._integration_ownership(tmp_path) assert github_owned == frozenset({test_path, github_path}) assert github_findings == () - workflow.write_text(yaml.safe_dump({"jobs": {}})) + per_pr: Final = tmp_path / ".github/workflows/test-mcp-interactions.yml" + per_pr.write_text(yaml.safe_dump({"jobs": {"mcp": {"steps": [{"run": f"pytest {github_path}"}]}}})) + _, shared_findings = coverage._integration_ownership(tmp_path) + assert [(finding.subject, finding.detail) for finding in shared_findings] == [ + (github_path, "GitHub-owned integration contract runs outside a dedicated integration workflow") + ] + per_pr.unlink() + dedicated.write_text(yaml.safe_dump({"jobs": {}})) _, missing_invocation = coverage._integration_ownership(tmp_path) assert [(finding.subject, finding.detail) for finding in missing_invocation] == [ (github_path, "GitHub-owned integration contract has no invoking job") @@ -94,6 +106,72 @@ def test_integration_groups_require_exclusive_scheduled_circleci_owner(tmp_path: assert circle_invocation == () +def _write_integration_tree(root: Path, github_path: str) -> None: + contracts: Final = root / "tests/integration/management" + contracts.mkdir(parents=True) + (contracts / "test_contract.py").write_text("def test_contract(): pass\n") + (root / github_path).parent.mkdir(parents=True, exist_ok=True) + (root / github_path).write_text("def test_contract(): pass\n") + (root / "tests/integration/run.py").write_text( + "from types import MappingProxyType\nfrom typing import Final\n" + 'GROUPS: Final = MappingProxyType({"management": ("management",)})\n' + f"GITHUB_FILES: Final = frozenset({{{github_path!r}}})\n" + ) + (root / ".circleci").mkdir() + (root / ".circleci/config.yml").write_text( + yaml.safe_dump( + { + "jobs": { + "integration_contracts": { + "steps": [{"run": {"command": "bash .circleci/scripts/run_integration.sh management"}}] + } + }, + "workflows": {"integration": {"jobs": [{"integration_contracts": {"suite": "management"}}]}}, + } + ) + ) + (root / ".github/workflows").mkdir(parents=True) + + +@pytest.mark.parametrize( + ("github_path", "workflow_name", "exit_code"), + [ + ("tests/integration/database/image/test_github_contract.py", "image-scan.yml", 0), + ("tests/integration/database/image/test_github_contract.py", "test-unit.yml", 1), + ("tests/integration/management/test_github_contract.py", "image-scan.yml", 1), + ], +) +def test_census_fails_when_a_github_owned_contract_runs_outside_its_dedicated_workflow( + tmp_path: Path, github_path: str, workflow_name: str, exit_code: int +) -> None: + script: Final = tmp_path / ".github/scripts/assert_ci_coverage.py" + script.parent.mkdir(parents=True) + shutil.copy(_MODULE_PATH, script) + _write_integration_tree(tmp_path, github_path) + (tmp_path / ".github/workflows" / workflow_name).write_text( + yaml.safe_dump({"jobs": {"tests": {"steps": [{"run": f"pytest {github_path}"}]}}}) + ) + result: Final = subprocess.run([sys.executable, "-I", str(script)], capture_output=True, text=True, check=False) + assert result.returncode == exit_code, result.stdout + assert "Playwright" not in result.stdout + + +def test_an_integration_file_outside_every_group_cannot_run_on_a_github_shard(tmp_path: Path) -> None: + _write_integration_tree(tmp_path, "tests/integration/database/image/test_github_contract.py") + ungrouped: Final = "tests/integration/ci_probe/test_ci_probe.py" + (tmp_path / ungrouped).parent.mkdir(parents=True) + (tmp_path / ungrouped).write_text("def test_probe(): pass\n") + _, before = coverage._integration_ownership(tmp_path) + (tmp_path / ".github/workflows/test-unit.yml").write_text( + yaml.safe_dump({"jobs": {"mcp": {"steps": [{"run": f"pytest {ungrouped}"}]}}}) + ) + _, after = coverage._integration_ownership(tmp_path) + assert [(finding.subject, finding.detail) for finding in after if finding.subject == ungrouped] == [ + (ungrouped, "integration file outside GITHUB_FILES is selected by GitHub Actions") + ] + assert not [finding for finding in before if finding.subject == ungrouped] + + def test_circleci_postgres_suites_count_their_test_path_as_invoked() -> None: config: Final = { "workflows": { diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cc2a1d62d52..c61ebb3737c 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -6426,6 +6426,78 @@ def test_whisper_azure(): assert round(cost, 5) == round(expected_cost, 5) +@pytest.mark.parametrize( + ("provider", "requested_model", "response_extras", "expected_seconds"), + [ + ("openai", "openai/gpt-fake-transcribe", {}, 2.0), + ("openai", "openai/gpt-fake-transcribe", {"_audio_transcription_duration": 1.35}, 1.35), + ("openai", "openai/gpt-fake-transcribe", {"duration": 1.34}, 1.34), + ("openai", "openai/whisper-fake", {}, 0.0), + ("azure", "gpt-fake-transcribe", {}, 0.0), + ], + ids=[ + "undecodable_upload_bills_provider_seconds", + "measured_duration_still_wins", + "provider_duration_still_wins", + "whisper_models_never_bill_provider_seconds", + "azure_never_bills_provider_seconds", + ], +) +def test_transcription_cost_falls_back_to_provider_billed_seconds( + monkeypatch: pytest.MonkeyPatch, + provider: str, + requested_model: str, + response_extras: Mapping[str, float], + expected_seconds: float, +) -> None: + model: Final = requested_model.removeprefix(f"{provider}/") + monkeypatch.setitem( + litellm.model_cost, + model, + {"input_cost_per_second": 0.01, "litellm_provider": provider, "mode": "audio_transcription"}, + ) + transcription: Final = convert_to_model_response_object( + response_object={ + "text": "What is the weather in Paris?", + "usage": {"type": "duration", "seconds": 2}, + **response_extras, + }, + model_response_object=TranscriptionResponse(), + hidden_params={"model": model, "custom_llm_provider": provider}, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, model=requested_model, call_type="atranscription" + ) + + assert cost == pytest.approx(0.01 * expected_seconds) + + +def test_custom_priced_transcription_deployment_bills_provider_seconds(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "priced-deployment", + {"input_cost_per_second": 0.000075, "litellm_provider": "openai", "mode": "audio_transcription"}, + ) + transcription: Final = convert_to_model_response_object( + response_object={"text": "What is the weather in Paris?", "usage": {"type": "duration", "seconds": 2}}, + model_response_object=TranscriptionResponse(), + hidden_params={"model": "gpt-fake-transcribe", "custom_llm_provider": "openai"}, + response_type="audio_transcription", + ) + + cost: Final = completion_cost( + completion_response=transcription, + model="openai/gpt-fake-transcribe", + call_type="atranscription", + custom_pricing=True, + router_model_id="priced-deployment", + ) + + assert cost == pytest.approx(0.000075 * 2) + + def test_gpt_image_2_azure_cost_tracking(): azure_image_generation_response: Final = { "created": 1758585600, diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index b12d3b34c8b..ea484c4a2f4 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -419,6 +419,24 @@ def test_azure_gpt_6_1_sol_data_zone_rows_price_requests_at_the_regional_premium assert total_cost(f"azure/{region}/gpt-6.1-sol") == pytest.approx(global_cost * premium) +AZURE_GPT_6_1_SOL_ROWS: Final = ( + "azure/gpt-6.1-sol", + "azure/gpt-6.1-sol-2026-09-29", + "azure/us/gpt-6.1-sol", + "azure/eu/gpt-6.1-sol", + "azure/apac/gpt-6.1-sol", + "azure_ai/gpt-6.1-sol", +) + + +@pytest.mark.parametrize("row", AZURE_GPT_6_1_SOL_ROWS) +def test_azure_gpt_6_1_sol_rows_accept_the_same_reasoning_efforts_as_the_openai_row(prices: dict, row: str): + openai_efforts: Final = resolve_supported_reasoning_efforts(prices["gpt-6.1-sol"], deployment_is_mapped=True) + assert openai_efforts + + assert resolve_supported_reasoning_efforts(prices[row], deployment_is_mapped=True) == openai_efforts + + # https://ai.google.dev/gemini-api/docs/models/deep-research-preview-04-2026 @pytest.mark.parametrize("model", ["gemini/deep-research-preview-04-2026", "gemini/deep-research-max-preview-04-2026"]) def test_gemini_deep_research_previews_accept_the_documented_input_window(prices: dict, model: str): diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index a5d74092d16..e73b9950594 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -4890,6 +4890,14 @@ LLMS_FORWARDER_CASES: Final = ( "staticmethod", False, ), + ( + "litellm.llms.openai.common_utils", + "BaseOpenAILLM", + "_get_sync_http_client", + "get_sync_http_client", + "staticmethod", + False, + ), ( "litellm.llms.openai.openai", "OpenAIChatCompletion", diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 6bc52ab3dfc..50a9d98385d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -38,6 +38,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.plumbing.context import set_request_root_span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.llms.anthropic.pass_through.messages.streaming_iterator import closing_superseded_attempt from litellm.llms.anthropic.pass_through.messages.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) @@ -9739,6 +9740,96 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice assert kwargs["tool_choice"] == "none" +@pytest.mark.parametrize( + "deployment_params, router_defaults, request_params, expected_tools", + [ + ({"drop_params": True}, {}, {}, [{"type": "web_search"}]), + ({}, {"drop_params": True}, {}, [{"type": "web_search"}]), + ({}, {}, {"additional_drop_params": ["tools"]}, [{"type": "web_search"}]), + ({"drop_params": True}, {}, {"drop_params": False}, 5), + ({}, {}, {}, 5), + ], + ids=[ + "deployment_drop_params", + "router_default_drop_params", + "request_drops_tools", + "request_overrides_deployment_drop", + "no_drop_policy", + ], +) +def test_update_kwargs_with_deployment_keeps_deployment_tools_when_malformed_request_tools_drop( + deployment_params: dict[str, object], + router_defaults: dict[str, object], + request_params: dict[str, object], + expected_tools: object, +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "o3-deep-research", + "litellm_params": { + "model": "openai/o3-deep-research", + "api_key": "fake-key", + "tools": [{"type": "web_search"}], + **deployment_params, + }, + }, + ], + default_litellm_params={**router_defaults}, + ) + kwargs: Final[dict[str, object]] = {"metadata": {}, "tools": 5, **request_params} + router._update_kwargs_with_deployment( + deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs + ) + + assert kwargs["tools"] == expected_tools + + +_NAMED_REQUEST_TOOL_CHOICE: Final = {"type": "function", "function": {"name": "lookup_order"}} + + +@pytest.mark.parametrize( + "deployment_params, expected_tool_fields", + [ + ( + {"drop_params": True, "tool_choice": "required"}, + {"tools": [{"type": "web_search"}], "tool_choice": "required"}, + ), + ({"drop_params": True}, {"tools": [{"type": "web_search"}]}), + ({}, {"tools": 5, "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, "parallel_tool_calls": False}), + ], + ids=["deployment_tool_choice_applies", "no_deployment_tool_choice", "no_drop_policy"], +) +def test_update_kwargs_with_deployment_drops_request_tool_settings_with_malformed_request_tools( + deployment_params: dict[str, object], expected_tool_fields: dict[str, object] +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "o3-deep-research", + "litellm_params": { + "model": "openai/o3-deep-research", + "api_key": "fake-key", + "tools": [{"type": "web_search"}], + **deployment_params, + }, + }, + ], + ) + kwargs: Final[dict[str, object]] = { + "metadata": {}, + "tools": 5, + "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, + "parallel_tool_calls": False, + } + router._update_kwargs_with_deployment( + deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs + ) + + tool_fields: Final = ("tools", "tool_choice", "parallel_tool_calls") + assert {name: kwargs[name] for name in tool_fields if name in kwargs} == expected_tool_fields + + def test_credential_name_injected_as_tag(): """ Test that litellm_credential_name from deployment litellm_params @@ -16004,6 +16095,7 @@ class _AnthropicMessagesFakeByteStream: self._chunks = list(chunks) self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} self.closed = False + self.closed_as_superseded = False def __aiter__(self): return self @@ -16014,6 +16106,7 @@ class _AnthropicMessagesFakeByteStream: return self._chunks.pop(0) async def aclose(self) -> None: + self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() self.closed = True @@ -17256,6 +17349,30 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +@pytest.mark.asyncio +async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): + """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream + logging must be queued before the retry starts producing its own: the router closes the dropped stream + before it asks the group for another.""" + router = _anthropic_messages_retry_router(num_retries=1) + superseded = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + ) + superseded_closed_when_retry_opened: list[bool] = [] + + def retry_stream(): + superseded_closed_when_retry_opened.append(superseded.closed_as_superseded) + return _anthropic_messages_retried_stream() + + provider = _AnthropicMessagesScriptedProvider(lambda: superseded, retry_stream) + + stream = await _anthropic_messages_stream_through_router(router, provider) + body = [chunk async for chunk in stream] + + assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] + assert superseded_closed_when_retry_opened == [True] + + @pytest.mark.asyncio async def test_a_ptu_share_holders_messages_stream_dropped_before_content_is_retried_on_its_shared_group(monkeypatch): """The same-group retry has to look the group's health up as the calling team; looked up with no @@ -18387,7 +18504,7 @@ async def test_aanthropic_messages_recover_stream_error_flushes_buffered_frames_ with patch.object(router, "_aanthropic_messages_fallback_attempt") as mock_attempt: recovery = router._aanthropic_messages_recover_stream_error( - declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper() + declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() ) with pytest.raises(BedrockError) as exc_info: await drain(recovery) @@ -18408,7 +18525,7 @@ async def test_aanthropic_messages_recover_stream_error_hands_converted_raise_to with patch.object(router, "_aanthropic_messages_fallback_attempt", new=fake_attempt): recovery = router._aanthropic_messages_recover_stream_error( - raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper() + raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() ) collected = [chunk async for chunk in recovery] assert collected == [b"fallback"] @@ -18466,6 +18583,25 @@ async def test_anthropic_messages_post_first_chunk_error_skips_fallback(): mock_fallback.assert_not_awaited() +@pytest.mark.asyncio +async def test_anthropic_messages_stream_the_client_closes_at_a_forwarded_error_is_not_closed_as_superseded(): + """An error frame forwarded to the client ends that request, and the SDK hangs up on it. No retry took the + request over, so the source must not be closed as a superseded attempt, which would log it as a failure and + drop the usage the provider billed for it.""" + router = _anthropic_messages_make_router() + error_chunk = _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("partial answer"), error_chunk]) + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + async for chunk in wrapped: + if chunk == error_chunk: + break + await wrapped.aclose() + + assert source.closed is True + assert source.closed_as_superseded is False + + @pytest.mark.asyncio async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames(): """A non-retriable error arriving while lifecycle frames are still diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 98025f3a79a..632ab020409 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2418,6 +2418,87 @@ def test_register_model_openrouter_without_slash(): assert "openai/gpt-4-turbo" in litellm.openrouter_models +def test_register_model_makes_the_model_resolvable_and_wildcard_visible_without_a_restart(): + model_name: Final = "my-tuned-gemini-regression" + try: + litellm.register_model( + { + model_name: { + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + } + }, + persist_across_reloads=False, + ) + + _, provider, _, _ = litellm.get_llm_provider(model=model_name) + + assert provider == "vertex_ai" + assert model_name in litellm.models_by_provider["vertex_ai"] + finally: + litellm.model_cost.pop(model_name, None) + litellm.vertex_language_models.discard(model_name) + litellm.models_by_provider["vertex_ai"].discard(model_name) + + +@pytest.mark.parametrize( + ("litellm_provider", "wildcard", "provider_set"), + [ + ("ai21", "ai21/*", "ai21_chat_models"), + ("vertex_ai-code-chat-models", "vertex_ai/*", "vertex_code_chat_models"), + ], +) +def test_registered_chat_model_lists_under_its_provider_wildcard( + litellm_provider: str, wildcard: str, provider_set: str +) -> None: + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model_name: Final = f"chat-listing-regression-{litellm_provider}" + try: + litellm.register_model( + {model_name: {"litellm_provider": litellm_provider, "mode": "chat"}}, persist_across_reloads=False + ) + + listed: Final = get_known_models_from_wildcard(wildcard) + + assert f"{wildcard.removesuffix('*')}{model_name}" in listed + finally: + litellm.model_cost.pop(model_name, None) + getattr(litellm, provider_set).discard(model_name) + litellm.add_known_models(model_cost_map={}) + + +def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing(): + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + deployment_id: Final = "priced-deployment-id-regression" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "priced-gemini", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "fake-key", + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + }, + "model_info": {"id": deployment_id, "litellm_provider": "gemini", "mode": "chat"}, + }, + {"model_name": "gemini/*", "litellm_params": {"model": "gemini/*", "api_key": "fake-key"}}, + ] + ) + + router.completion(model="priced-gemini", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + listed: Final = get_known_models_from_wildcard("gemini/*") + + assert "gemini/gemini-2.5-flash" in listed + assert f"gemini/{deployment_id}" not in listed + assert deployment_id not in litellm.models_by_provider["gemini"] + assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.000001 + + def test_reasoning_content_preserved_in_text_completion_wrapper(): """Ensure reasoning_content is copied from delta to text_choices.""" chunk = ModelResponseStream( @@ -9585,6 +9666,33 @@ def test_get_provider_audio_transcription_config_resolves_for_every_provider() - assert isinstance(configs[LlmProviders.OPENAI], litellm.OpenAIWhisperAudioTranscriptionConfig) +@pytest.mark.parametrize( + ("model", "optional_params", "expected_response_format"), + [ + ("whisper-1", {}, "verbose_json"), + ("my-transcribe-whisper", {}, "verbose_json"), + ("whisper-1", {"response_format": "srt"}, "srt"), + ("gpt-4o-mini-transcribe", {}, None), + ("gpt-transcribe", {}, None), + ("gpt-transcribe", {"response_format": "text"}, "text"), + ], +) +def test_openai_transcription_defaults_verbose_json_only_for_whisper_models( + model: str, optional_params: dict[str, str], expected_response_format: str | None +) -> None: + config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=model, provider=LlmProviders.OPENAI + ) + assert config is not None + + request: Final = config.transform_audio_transcription_request( + model=model, audio_file=b"audio", optional_params=dict(optional_params), litellm_params={} + ) + + assert isinstance(request.data, dict) + assert request.data.get("response_format") == expected_response_format + + def test_get_valid_models_from_provider_cache_invalidation(monkeypatch): """ Test that get_valid_models returns the correct models for a given provider diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 122d811e491..aef43a15a72 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -65,6 +65,7 @@ CONNECTION_NAMES: Final = ( "api_key", "api_base", "api_version", + "responses_api_version", "region_name", "headers", "provider_specific_header", diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 6188bb966ee..0db122b927f 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -63,6 +63,7 @@ def test_custom_pricing_params_keeps_every_field_it_had(): "output_cost_per_character", "cache_read_input_token_cost", "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", "cost_per_second", "input_cost_per_second", "cache_read_input_token_cost_flex", @@ -372,3 +373,9 @@ def test_server_owned_wif_fields_present_reports_openai_fields(): def test_reject_server_owned_wif_params_names_each_openai_field(param_name: str): with pytest.raises(ValueError, match=param_name): reject_server_owned_wif_params({param_name: "client-supplied"}) + + +@pytest.mark.parametrize("param_name", openai_wif_litellm_params) +def test_reject_server_owned_wif_params_catches_the_array_form_key(param_name: str): + with pytest.raises(ValueError, match=param_name): + reject_server_owned_wif_params({f"{param_name}[]": "client-supplied"}) diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index f6c888d20ff..83b59d77aaf 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -349,6 +349,56 @@ async def test_malformed_vector_store_index_row_from_the_db_raises_a_validation_ await VectorStoreIndexRegistry.get_vector_store_indexes_from_db(_database_listing_index_rows([row])) +@pytest.mark.parametrize( + ("tools", "expected"), + [ + (5, []), + (["vector_store_ids", {"type": "file_search", "vector_store_ids": ["vs_kept"]}], ["vs_kept"]), + ( + [{"type": "file_search", "vector_store_ids": 5}, {"type": "file_search", "vector_store_ids": ["vs_kept"]}], + ["vs_kept"], + ), + ], + ids=["int", "string_entry", "int_vector_store_ids"], +) +def test_vector_store_ids_to_run_skips_malformed_tools(tools: object, expected: list[str]) -> None: + assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == expected + + +@pytest.mark.parametrize( + "vector_store_ids", + ["vs_requested", ["vs_requested", 5]], + ids=["string", "list_with_a_non_string"], +) +def test_vector_store_ids_to_run_reads_every_store_id_a_malformed_vector_store_ids_field_names( + vector_store_ids: object, +) -> None: + tools: Final = [{"type": "file_search", "vector_store_ids": vector_store_ids}] + assert VectorStoreRegistry().get_vector_store_ids_to_run(non_default_params={}, tools=tools) == ["vs_requested"] + + +@pytest.mark.parametrize( + "vector_store_ids", + ["vs_allowed", ["vs_allowed", 12345], [12345], ["vs_allowed"]], + ids=["string", "list_with_a_non_string", "only_non_strings", "list_of_strings"], +) +def test_vector_store_tools_pop_exactly_the_store_ids_the_registry_reports_to_run(vector_store_ids: object) -> None: + now: Final = datetime.now(timezone.utc) + registry: Final = VectorStoreRegistry( + [ + LiteLLM_ManagedVectorStore(vector_store_id=vector_store_id, created_at=now, updated_at=now) + for vector_store_id in ("vs_allowed", 12345) + ] + ) + tools: Final = [{"type": "file_search", "vector_store_ids": vector_store_ids}] + checked: Final = registry.get_vector_store_ids_to_run(non_default_params={}, tools=tools) + ran: Final[list[str]] = [] + + registry.get_and_pop_recognised_vector_store_tools(tools=[dict(tool) for tool in tools], vector_store_ids=ran) + + assert ran == checked + + class TestRegistriesDoNotShareDefaultState: def test_adding_to_one_registry_leaves_the_next_one_empty(self): VectorStoreRegistry().add_vector_store_to_registry( diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index b1b6cef7320..9eead41a102 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2090,11 +2090,26 @@ "count": 1 } }, + "src/components/ui/attachment.tsx": { + "local/no-ad-hoc-z-index": { + "count": 2 + } + }, + "src/components/ui/message-scroller.tsx": { + "local/no-noop-hover-variant": { + "count": 1 + } + }, "src/components/ui/slider.tsx": { "no-nested-ternary": { "count": 1 } }, + "src/components/ui/toggle-group.tsx": { + "local/no-ad-hoc-z-index": { + "count": 2 + } + }, "src/components/update_model_credentials_modal.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx index 3eb80d3f587..e6cd9eb54a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/PassthroughAuthorizeSection.tsx @@ -3,9 +3,11 @@ import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { Label } from "@/components/ui/label"; import { PasswordInput } from "@/components/shared/PasswordInput"; +import { Input } from "@/components/ui/input"; +import { MultiSelect } from "@/components/shared/MultiSelect"; import DcrBridgeToggle from "./DcrBridgeToggle"; import { MountedFormField } from "@/components/common_components/MountedFormField"; -import { textControl } from "./mcpFieldRules"; +import { tagsControl, textControl } from "./mcpFieldRules"; import { credentialAuthClass, isClientForwardedTokenMode } from "@/components/mcp_tools/types"; interface PassthroughOAuthFlow { @@ -69,6 +71,7 @@ export default function PassthroughAuthorizeSection({ const clientSecretPlaceholder = classUnchanged ? "Leave blank to keep the currently saved secret (if any)" : "Leave blank for public clients / PKCE"; + const fieldClassName = "rounded-lg border-border focus:border-info focus:ring-ring"; const clientIdExtra = classUnchanged ? "Set this to make everyone authorize through a specific app; required for upstreams without dynamic client registration (e.g. a pre-registered Slack app)." : "Switching the auth type discards the previously saved app; enter a client ID here or leave blank to use dynamic client registration."; @@ -113,6 +116,39 @@ export default function PassthroughAuthorizeSection({ /> )} + Authorization URL (optional)} + name="authorization_url" + help="Set this when the upstream publishes no OAuth metadata, for example an internal server behind Microsoft Entra. Leave blank to discover it from the upstream." + > + {(control) => ( + + )} + + Token URL (optional)} + name="token_url" + help="Token endpoint of the same authorization server. Leave blank to discover it from the upstream." + > + {(control) => ( + + )} + + Scopes (optional)} + name={["credentials", "scopes"]} + help="Scopes requested from the upstream when a user authorizes. Press enter or comma between scopes." + > + {(control) => } + {isEditing && onRemoveStoredAppChange && (