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/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index f8a7c38ef79..636a5031e5f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -16,6 +16,10 @@ on: - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" + - "litellm/chat_completions/dispatch.py" + - "litellm/embeddings/dispatch.py" + - "litellm/messages/dispatch.py" + - "litellm/responses/dispatch.py" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - "Makefile" @@ -46,6 +50,10 @@ on: - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" + - "litellm/chat_completions/dispatch.py" + - "litellm/embeddings/dispatch.py" + - "litellm/messages/dispatch.py" + - "litellm/responses/dispatch.py" - "litellm/llms/base_llm/ocr/**" - "litellm/llms/custom_httpx/llm_http_handler.py" - "Makefile" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0558875bdad..1b486d78f04 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -696,7 +696,6 @@ jobs: tests/unit/responses/litellm_completion_transformation tests/unit/responses/test_additional_tools.py tests/unit/responses/test_custom_tool_call.py - tests/unit/responses/test_dispatch.py tests/unit/responses/test_metadata_codex_callback.py tests/unit/responses/test_no_duplicate_spend_logs.py tests/unit/responses/test_null_test_fix.py 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/README.md b/README.md index 3bd0fd54a65..77f7e3baa4d 100644 --- a/README.md +++ b/README.md @@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | | | [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | | | [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | | +| [Nace (`nace`)](https://console.nace.ai/docs/drex) | | | | | | | | | | | | [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | | | [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index c7a8e3e3294..d612a042e03 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [ "x-openrouter-experimental-metadata", ]; const ANTHROPIC_API_HEADER_PROVIDERS: &str = - "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter"; + "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek"; pub async fn create( State(gateway): State>, diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index b0a9baa2fac..bf11e67e0fc 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -16,7 +16,7 @@ use wiremock::{ #[rstest] #[case::anthropic("anthropic/test-model", "/v1/messages", true)] #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)] -#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)] +#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)] #[tokio::test] async fn messages_reaches_the_provider_and_preserves_json_or_sse( #[case] model: &str, diff --git a/litellm/__init__.py b/litellm/__init__.py index 88e56c14187..cdcf9d1d80c 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, @@ -1691,6 +1694,9 @@ if TYPE_CHECKING: from .llms.perplexity.decisions.transformation import ( PerplexityDecisionsConfig as PerplexityDecisionsConfig, ) + from .llms.nace.decisions.transformation import ( + NaceDecisionsConfig as NaceDecisionsConfig, + ) from .llms.typesafe.decisions.transformation import ( TypeSafeDecisionsConfig as TypeSafeDecisionsConfig, ) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index df549701cf4..244efc8047b 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -37,6 +37,16 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families" _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None) +@contextmanager +def internal_sub_call() -> Generator[None]: + """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own.""" + token: Final = is_internal_call.set(True) + try: + yield + finally: + is_internal_call.reset(token) + + @contextmanager def post_response_phase() -> Generator[None]: """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns.""" diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index d47d72ef7a8..8abc24d1417 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -162,6 +162,7 @@ LLM_CONFIG_NAMES: Final = ( "HostedVLLMRerankConfig", "PerplexityDecisionsConfig", "TypeSafeDecisionsConfig", + "NaceDecisionsConfig", "OpenRouterDecisionsConfig", "FireworksAIDecisionsConfig", "CloudflareDecisionsConfig", @@ -729,6 +730,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "HostedVLLMRerankConfig", ), "PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"), + "NaceDecisionsConfig": (".llms.nace.decisions.transformation", "NaceDecisionsConfig"), "TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"), "OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"), "FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"), 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/caching/caching_handler.py b/litellm/caching/caching_handler.py index 501dc266ba0..11488a03f3f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError import litellm -from litellm._internal_context import post_response_phase +from litellm._internal_context import is_internal_call, post_response_phase from litellm._logging import print_verbose, verbose_logger from litellm.caching import InMemoryCache from litellm.caching.caching import S3Cache, response_cache_phase @@ -511,7 +511,9 @@ class LLMCachingHandler: custom_llm_provider=custom_llm_provider, ) - if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result): + if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks( + cached_result=cached_result + ): logging_obj.handle_sync_success_callbacks_for_async_calls( result=cached_result, start_time=start_time, @@ -820,7 +822,12 @@ class LLMCachingHandler: start_time (datetime): The start time of the operation. end_time (datetime): The end time of the operation. cache_hit (bool): Whether it was a cache hit. + + An internal sub-call's hit fires nothing, since the parent call that folds it in logs it. """ + if is_internal_call.get(): + return + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( 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..f448fc351c3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,7 +391,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1" DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" -BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime" @@ -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 0e315e97432..fef6d618e9d 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"] @@ -1426,32 +1438,15 @@ def completion_cost( vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") ) -> float: """ - Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. + Return the USD cost of a LiteLLM call. - Parameters: - completion_response (litellm.ModelResponses): [Required] The response received from a LiteLLM completion request. + Pass `completion_response` to price a response LiteLLM returned. Without one, pass `model` plus + `messages` / `prompt` and `completion`, and token counts are computed from the text. Prices come from + `litellm.model_cost` unless `custom_cost_per_token` or `custom_cost_per_second` is given. `total_time` + is in milliseconds and only matters for per-second pricing such as Replicate and SageMaker. - [OPTIONAL PARAMS] - model (str): Optional. The name of the language model used in the completion calls - prompt (str): Optional. The input prompt passed to the llm - completion (str): Optional. The output completion text from the llm - total_time (float, int): Optional. (Only used for Replicate LLMs) The total time used for the request in seconds - custom_cost_per_token: Optional[CostPerToken]: the cost per input + output token for the llm api call. - custom_cost_per_second: Optional[float]: the cost per second for the llm api call. - - Returns: - float: The cost in USD dollars for the completion based on the provided parameters. - - Exceptions: - Raises exception if model not in the litellm model cost map. Register model, via custom pricing or PR - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json - - - Note: - - If completion_response is provided, the function extracts token information and the model name from it. - - If completion_response is not provided, the function calculates token counts based on the model and input text. - - The cost is calculated based on the model, prompt tokens, and completion tokens. - - For certain models containing "togethercomputer" in the name, prices are based on the model size. - - For un-mapped Replicate models, the cost is calculated based on the total time used for the request. + Raises if no candidate model name resolves to pricing. Register the model via `litellm.register_model` + or add it to model_prices_and_context_window.json. """ try: call_type = infer_call_type(call_type, completion_response) or "completion" @@ -1739,6 +1734,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): @@ -2048,7 +2049,7 @@ def completion_cost( ) if idx == len(potential_model_names) - 1: raise e - raise Exception(f"Unable to calculat cost for received potential model names - {potential_model_names}") + raise Exception(f"Unable to calculate cost for received potential model names - {potential_model_names}") except Exception as e: raise e @@ -2063,13 +2064,20 @@ def get_response_cost_from_hidden_params( additional_headers: Final = _hidden_params_dict.get("additional_headers", {}) if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers: - response_cost: Final = additional_headers["llm_provider-x-litellm-response-cost"] - if response_cost is None: - return None - return float(additional_headers["llm_provider-x-litellm-response-cost"]) + return _billable_header_cost(additional_headers["llm_provider-x-litellm-response-cost"]) return None +def _billable_header_cost(value: object) -> float | None: + if not isinstance(value, (str, int, float)): + return None + try: + cost: Final = float(value) + except ValueError: + return None + return cost if math.isfinite(cost) and cost >= 0 else None + + def response_cost_calculator( response_object: ModelResponse | EmbeddingResponse diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index ef35b61f10e..4943b22e1fb 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -57,7 +57,7 @@ from litellm.types.decisions import ( UnsupportedDecisionsRequest, ) from litellm.types.utils import LlmProviders, ProviderSpecificHeader -from litellm.utils import ProviderConfigManager +from litellm.utils import ProviderConfigManager, get_model_cost_key SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]] OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]] @@ -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}'", @@ -666,7 +667,10 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci { "model": f"{call.custom_llm_provider}/{call.model}", "custom_llm_provider": call.custom_llm_provider, - "provider_response_model": f"{call.custom_llm_provider}/{call.model}", + "provider_response_model": ( + get_model_cost_key(f"{call.custom_llm_provider}/{response.model}") + or f"{call.custom_llm_provider}/{call.model}" + ), **({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}), } ) 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/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 42f605443c7..86f80047e52 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -103,6 +103,43 @@ def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_ ) +def langfuse_trace_metadata(metadata: object, proxy_server_request: object) -> Mapping[str, object]: + entries: Final = _object_mapping(metadata) or _NO_METADATA + request: Final = _object_mapping(proxy_server_request) + headers: Final = _object_mapping(request.get("headers")) if request is not None else None + if headers is None: + return entries + header_entries: Final = { + key.removeprefix("langfuse_"): value + for key, value in _widened_items(headers) + if isinstance(key, str) and key.startswith("langfuse_") + } + return MappingProxyType({**entries, **header_entries}) + + +def langfuse_trace_id( + metadata: Mapping[str, object], + standard_trace_id: object, + litellm_trace_id: object, + litellm_call_id: object, + proxy_server_request: object, +) -> object: + existing_trace_id: Final = metadata.get("existing_trace_id") + if existing_trace_id is not None: + return existing_trace_id + metadata_trace_id: Final = metadata.get("trace_id") + trace_id: Final = ( + metadata_trace_id + if metadata_trace_id is not None + else standard_trace_id + if standard_trace_id is not None + else litellm_trace_id or litellm_call_id + ) + if _is_session_header_trace(trace_id, metadata.get("session_id"), proxy_server_request): + return litellm_call_id or trace_id + return trace_id + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -680,33 +717,19 @@ class LangFuseLogger: # Add default langfuse tags tags = self.add_default_langfuse_tags(tags=tags, kwargs=kwargs, metadata=metadata) + call_trace_id: Final = langfuse_trace_id( + metadata=clean_metadata, + standard_trace_id=standard_logging_object.get("trace_id") + if standard_logging_object is not None + else None, + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_call_id=litellm_call_id, + proxy_server_request=litellm_params.get("proxy_server_request"), + ) session_id: Final = clean_metadata.pop("session_id", None) trace_name = cast(str | None, clean_metadata.pop("trace_name", None)) - trace_id = clean_metadata.pop("trace_id", None) - # Use standard_logging_object.trace_id if available (when trace_id from metadata is None) - # This allows standard trace_id to be used when provided in standard_logging_object - if trace_id is None and standard_logging_object is not None: - trace_id = cast(str | None, standard_logging_object.get("trace_id")) - # Fallback to litellm_call_id if no trace_id found - if trace_id is None: - trace_id = kwargs.get("litellm_trace_id") or litellm_call_id + clean_metadata.pop("trace_id", None) existing_trace_id: Final = clean_metadata.pop("existing_trace_id", None) - # If existing_trace_id is provided, use it as the trace_id to return - # This allows continuing an existing trace while still returning the correct trace_id - if existing_trace_id is not None: - trace_id = existing_trace_id - call_trace_id: Final = ( - litellm_call_id or trace_id - if existing_trace_id is None - and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request")) - else trace_id - ) - if call_trace_id != trace_id: - verbose_logger.debug( - "Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace", - trace_id, - call_trace_id, - ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fa5496004e3..517a28cf44a 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 @@ -3368,7 +3373,7 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) - if remaining_requests: + if isinstance(remaining_requests, int): """ "model_group", "api_provider", @@ -3382,7 +3387,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) - if remaining_tokens: + if isinstance(remaining_tokens, int): _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index cc1f6fef019..e706bcadef8 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -318,7 +318,7 @@ def remove_items_at_indices(items: _PoppableByIndex | None, indices: Iterable[in items.pop(index) -def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: dict) -> dict: +def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metadata: Mapping[str, object]) -> dict: """ Helper to get litellm metadata for spend tracking @@ -326,7 +326,7 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad and user_api_key values are in 'metadata'. """ potential_spend_tracking_metadata_substring: Final = "user_api_key" - for key, value in metadata.items(): + for key, value in dict(metadata).items(): if potential_spend_tracking_metadata_substring in key: litellm_metadata[key] = value return litellm_metadata 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/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index 714b258413e..4cb9dcd53f1 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -77,6 +77,11 @@ def get_hidden_params(obj: object) -> MutableMapping[str, object] | None: return _as_hidden_params_mapping(hidden_params) +def served_from_cache(obj: object) -> bool: + hidden_params: Final = get_hidden_params(obj) + return hidden_params is not None and hidden_params.get("cache_hit") is True + + def set_hidden_params(obj: object, hidden_params: dict[str, object] | HiddenParams) -> None: if isinstance(obj, dict): obj[_HIDDEN_PARAMS_ATTR] = hidden_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fa009ca6327..35decbb4bd7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -79,7 +79,12 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_param +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + served_from_cache, + set_hidden_param, +) from litellm.litellm_core_utils.internal_call_metadata import ( MODEL_ACCESS_GROUP_METADATA_KEY, is_unbilled_non_inference_call, @@ -2056,7 +2061,7 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: if response_cost is None and not calculation_failed: return - if self.model_call_details.get("cache_hit") is True: + if self.model_call_details.get("cache_hit") is True or _served_from_cache_at_no_charge(result): self.model_call_details["zero_cost_diagnostic"] = None return try: @@ -2295,6 +2300,24 @@ 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.""" + usage: Final = self.model_call_details.pop("combined_usage_object", None) + response_cost: Final = self.model_call_details.pop("response_cost", None) + if isinstance(usage, Usage): + self.model_call_details["failed_attempts_usage"] = ( + *self.model_call_details.get("failed_attempts_usage", ()), + (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), + ) + def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" self.model_call_details["continued_partial_usage"] = usage @@ -2662,6 +2685,8 @@ class Logging(LiteLLMLoggingBaseClass): result = self._handle_a2a_response_logging(result=result) logging_result: Final = self.normalize_logging_result(result=result) + if cache_hit is None and _served_from_cache_at_no_charge(logging_result): + self.model_call_details["cache_hit"] = True if isinstance(result, Response) and isinstance( logging_result, (ModelResponse, EmbeddingResponse, TextCompletionResponse) @@ -4327,6 +4352,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 @@ -4427,6 +4455,7 @@ class Logging(LiteLLMLoggingBaseClass): """ from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, + responses_provider_metadata, ) served_id: Final = _provider_response_id(result) @@ -4451,7 +4480,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): @@ -6679,6 +6710,11 @@ def _get_status_fields( return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status) +def _served_from_cache_at_no_charge(result: object) -> bool: + hidden_params: Final = get_hidden_params(result) + return served_from_cache(result) and hidden_params is not None and hidden_params.get("response_cost") == 0 + + def _extract_response_obj_and_hidden_params( init_response_obj: object, original_exception: Exception | None, @@ -6841,8 +6877,9 @@ def get_standard_logging_object_payload( ) # maintain backwards compatibility with old request body check saved_cache_cost: float = 0.0 - if cache_hit is True: + if cache_hit is True or served_from_cache(init_response_obj): id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id + if cache_hit is True: saved_cache_cost = ( logging_obj.response_cost_calculator( result=init_response_obj, 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/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c9b96c0a4c2..927caec3350 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,17 +1,20 @@ import asyncio import json import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence from dataclasses import dataclass from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast, runtime_checkable from pydantic import TypeAdapter from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm._uuid import uuid from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -98,6 +101,16 @@ class ScopedWebSocket(Protocol): def scope(self) -> _ASGIScope: ... +@runtime_checkable +class _ReplayScopeWebSocket(Protocol): + @property + def scope(self) -> MutableMapping[str, object]: ... + + +def client_websocket_scope(websocket: object) -> MutableMapping[str, object] | None: + return websocket.scope if isinstance(websocket, _ReplayScopeWebSocket) else None + + class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... @@ -116,10 +129,87 @@ def _decode_json_object(payload: str) -> Mapping[str, object]: return json.loads(payload) +_GUARDRAIL_TURN_DETECTION_GATE: Final[Mapping[str, object]] = MappingProxyType( + {"type": "server_vad", "create_response": False} +) +_EMPTY_SESSION: Final[Mapping[str, object]] = MappingProxyType({}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_GUARDRAIL_EVENT_ID_PREFIX: Final = "litellm_guardrail_" + + +def _json_object(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, Mapping) else None + + +def _mapping_field(container: Mapping[str, object], key: str) -> Mapping[str, object] | None: + return _json_object(container.get(key)) + + +def _guardrail_event_id() -> str: + return f"{_GUARDRAIL_EVENT_ID_PREFIX}{uuid.uuid4().hex}" + + +def _error_event_id(event: Mapping[str, object]) -> str | None: + error: Final = _mapping_field(event, "error") + caused_by: Final = None if error is None else error.get("event_id") + return caused_by if isinstance(caused_by, str) else None + + +_TurnDetectionGate: TypeAlias = Callable[[object], Mapping[str, object]] + + +def _turn_detection_with_auto_response_off(turn_detection: object) -> Mapping[str, object]: + client_turn_detection: Final = _json_object(turn_detection) + if client_turn_detection is None: + return dict(_GUARDRAIL_TURN_DETECTION_GATE) + return {**client_turn_detection, "create_response": False} + + +def _nested_turn_detection_part(session: Mapping[str, object], gate: _TurnDetectionGate) -> Mapping[str, object]: + audio: Final = _mapping_field(session, "audio") + if audio is None: + return {} + audio_input: Final = _mapping_field(audio, "input") + if audio_input is None or "turn_detection" not in audio_input: + return {} + gated_input: Final = { + **audio_input, + "turn_detection": gate(audio_input["turn_detection"]), + } + return {"audio": {**audio, "input": gated_input}} + + +def _with_create_response_disabled( + session: Mapping[str, object], gate: _TurnDetectionGate +) -> Mapping[str, object] | None: + """``session`` with every turn_detection the client sent unable to auto-respond, or None when it sent none.""" + flat_part: Final = {"turn_detection": gate(session["turn_detection"])} if "turn_detection" in session else {} + nested_part: Final = _nested_turn_detection_part(session, gate) + if not flat_part and not nested_part: + return None + return {**session, **flat_part, **nested_part} + + +def _gate_session_update( + session: object, gate: _TurnDetectionGate, inject_gate_when_absent: bool +) -> Mapping[str, object] | None: + """The session to forward under a transcript guardrail, or None when the frame needs no change.""" + client_session: Final = _json_object(session) + if client_session is None: + return None + gated: Final = _with_create_response_disabled(client_session, gate) + if gated is not None or not inject_gate_when_absent: + return gated + return {**client_session, "turn_detection": gate(None)} + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... def patch_outgoing_session(self, session: dict) -> dict: ... + def gate_turn_detection(self, turn_detection: object, /) -> Mapping[str, object]: ... + @property + def deletes_conversation_items(self) -> bool: ... class RealtimeUsageMeter(Protocol): @@ -154,6 +244,7 @@ class RealTimeStreaming: logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket + self._replay_scope: Final = client_websocket_scope(websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker @@ -197,6 +288,8 @@ class RealTimeStreaming: # Track whether we have already sent the guardrail turn-detection update # that disables provider auto-response for transcription guardrails. self._guardrail_turn_detection_update_sent: bool = False + self._session_gate_sent: bool = False + self._guardrail_cancel_event_ids: frozenset[str] = frozenset() # Deferred Gemini Live setup: Pipecat may stream audio before session.update. # Buffer client audio until the backend acknowledges setup (setupComplete). self._backend_setup_complete: bool = provider_config is None or provider_config.requires_session_configuration() @@ -213,6 +306,12 @@ class RealTimeStreaming: self._is_transcription_session: bool = force_transcription_model is not None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer + self._gate_turn_detection: Final[_TurnDetectionGate] = ( + _turn_detection_with_auto_response_off if event_normalizer is None else event_normalizer.gate_turn_detection + ) + self._deletes_conversation_items: Final = ( + event_normalizer is None or event_normalizer.deletes_conversation_items + ) self._usage_meter = usage_meter @property @@ -734,20 +833,37 @@ class RealTimeStreaming: except (json.JSONDecodeError, TypeError): return + def _make_turn_detection_update(self, turn_detection: Mapping[str, object]) -> dict[str, object]: + session: Final[dict[str, object]] = ( + {"turn_detection": turn_detection} + if self._backend_uses_beta_protocol + else {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + ) + return {"type": "session.update", "session": session} + def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, str | bool]] = { - "type": "server_vad", - "create_response": False, - } - if self._backend_uses_beta_protocol: - session: dict[str, object] = {"turn_detection": turn_detection} - else: - session = { - "type": "realtime", - "audio": {"input": {"turn_detection": turn_detection}}, - } - return json.dumps({"type": "session.update", "session": session}) + return json.dumps(self._make_turn_detection_update(self._gate_turn_detection(None))) + + async def _send_guardrail_cancel(self) -> None: + event_id: Final = _guardrail_event_id() + self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id} + await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id})) + + async def _withdraw_blocked_turn(self, item_id: str | None) -> None: + await self._send_guardrail_cancel() + if item_id is not None and self._deletes_conversation_items: + await self._send_to_backend(json.dumps({"type": "conversation.item.delete", "item_id": item_id})) + + def _is_guardrail_cancel_error(self, event: Mapping[str, object]) -> bool: + return _error_event_id(event) in self._guardrail_cancel_event_ids + + async def _send_session_gate_once(self) -> None: + """Gate the backend's default turn_detection ahead of the first client session.update, which overrides it.""" + if self._session_gate_sent or not self._should_disable_vad_auto_response(): + return + self._session_gate_sent = True + await self._send_to_backend(self._make_disable_auto_response_message()) async def _maybe_send_guardrail_turn_detection_update(self) -> None: """Disable provider auto-response once when transcription guardrails are enabled.""" @@ -921,7 +1037,7 @@ class RealTimeStreaming: if pre_block_backend_message is not None: await self._send_to_backend(pre_block_backend_message) if not self._is_transcription_session: - await self._send_to_backend(json.dumps({"type": "response.cancel"})) + await self._withdraw_blocked_turn(item_id) # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( json.dumps( @@ -1051,10 +1167,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails( - cast(str, transcript), - item_id=cast(str | None, event.get("item_id")), - ) + blocked = await self.run_realtime_guardrails(cast(str, transcript)) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -1086,7 +1199,7 @@ class RealTimeStreaming: if event_type == "session.created" and self._should_disable_vad_auto_response(): self.store_message(event_obj) await self.websocket.send_text(self._event_to_client_json(event_obj)) - await self._send_to_backend(self._make_disable_auto_response_message()) + await self._send_session_gate_once() return True if event_type == "conversation.item.input_audio_transcription.completed": @@ -1140,6 +1253,10 @@ class RealTimeStreaming: if self._usage_meter is not None: self._usage_meter.observe_backend_event(event) + if event.get("type") == "error" and self._is_guardrail_cancel_error(event): + verbose_logger.debug("Realtime: not forwarding an error caused by a guardrail event: %s", event) + continue + if await self._handle_raw_backend_message(event, raw_response): continue @@ -1346,6 +1463,13 @@ class RealTimeStreaming: return item async def _receive_client_message(self) -> str | None: + refused_opening_message: Final = ( + None + if self._replay_scope is None + else self._replay_scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + ) + if isinstance(refused_opening_message, str): + return refused_opening_message try: return await self.websocket.receive_text() except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone @@ -1472,19 +1596,21 @@ class RealTimeStreaming: and self.session_configuration_request is None and not self._guardrail_turn_detection_update_sent and self._should_disable_vad_auto_response() - ): - session: Mapping[str, object] | None = msg_obj.setdefault("session", {}) - if isinstance(session, dict): - existing_td = session.get("turn_detection") - if not isinstance(existing_td, dict): - existing_td = {} - existing_td["create_response"] = False - session["turn_detection"] = existing_td - message = json.dumps(msg_obj) - guardrail_turn_detection_injected = True - verbose_logger.debug( - "Injected turn_detection into first session.update for audio transcription guardrails" + and ( + gated_first := _gate_session_update( + client_event.get("session", _EMPTY_SESSION), + self._gate_turn_detection, + inject_gate_when_absent=True, ) + ) + is not None + ): + msg_obj["session"] = gated_first + message = json.dumps(msg_obj) + guardrail_turn_detection_injected = True + verbose_logger.debug( + "Injected turn_detection into first session.update for audio transcription guardrails" + ) ## GUARDRAIL: Force ``create_response`` to False in any # client-provided ``turn_detection`` so a later @@ -1492,52 +1618,25 @@ class RealTimeStreaming: # and bypass the transcription guardrail after the # initial disable. Covers both the flat beta key and the # nested GA ``audio.input.turn_detection`` shape, since - # the GA remap below also accepts either form. Skipped - # when the injection block above already ran for this - # message, to avoid redundant double-serialization. + # the GA remap below also accepts either form. A later + # update with no turn_detection is forwarded as is: OpenAI + # keeps the earlier disable for an unset field, and xAI + # falls back to its manual default, which never auto-responds. if ( msg_type == "session.update" and not guardrail_turn_detection_injected and self._should_disable_vad_auto_response() + and ( + gated_later := _gate_session_update( + client_event.get("session"), + self._gate_turn_detection, + inject_gate_when_absent=False, + ) + ) + is not None ): - session = client_event.get("session") - if isinstance(session, dict): - td_overridden = False - flat_td = session.get("turn_detection") - flat_td_present = flat_td is not None - if flat_td_present: - if not isinstance(flat_td, dict): - flat_td = {} - if flat_td.get("create_response") is not False: - flat_td["create_response"] = False - session["turn_detection"] = flat_td - td_overridden = True - nested_td_present = False - audio = session.get("audio") - if isinstance(audio, dict): - audio_input = audio.get("input") - if isinstance(audio_input, dict): - nested_td = audio_input.get("turn_detection") - if nested_td is not None: - nested_td_present = True - if not isinstance(nested_td, dict): - nested_td = {} - if nested_td.get("create_response") is not False: - nested_td["create_response"] = False - audio_input["turn_detection"] = nested_td - td_overridden = True - # Symmetric with the first-update injection block: - # if the client omitted turn_detection entirely on - # a subsequent session.update, still inject the - # ``create_response: False`` override so the - # transcription guardrail cannot be re-enabled by - # any downstream merge that drops the original - # disable. - if not flat_td_present and not nested_td_present: - session["turn_detection"] = {"create_response": False} - td_overridden = True - if td_overridden: - message = json.dumps(msg_obj) + msg_obj["session"] = gated_later + message = json.dumps(msg_obj) # GA compatibility: remap beta-style session fields only when # the upstream is in GA mode. Beta upstreams expect the flat @@ -1590,6 +1689,8 @@ class RealTimeStreaming: # would permanently disable the injection if ``_send_to_backend`` # raised — neither this loop nor # ``_maybe_send_guardrail_turn_detection_update`` would retry. + if msg_type == "session.update" and self.provider_config is None: + await self._send_session_gate_once() sent = await self._send_to_backend(message) if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 95d062a190c..78fb82c98e7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1128,32 +1128,24 @@ class CustomStreamWrapper: """ - if self.merge_reasoning_content_in_choices is True: - reasoning_content: Final = getattr(model_response.choices[0].delta, "reasoning_content", None) - if reasoning_content: - self._emitted_disqualifying_content = True - if self.sent_first_thinking_block is False: - # Ensure content is not None before concatenation - if model_response.choices[0].delta.content is None: - model_response.choices[0].delta.content = "" - model_response.choices[0].delta.content += "" + reasoning_content - self.sent_first_thinking_block = True - elif ( - self.sent_first_thinking_block is True - and hasattr(model_response.choices[0].delta, "reasoning_content") - and model_response.choices[0].delta.reasoning_content - ): - model_response.choices[0].delta.content = reasoning_content - elif ( - self.sent_first_thinking_block is True - and not self.sent_last_thinking_block - and model_response.choices[0].delta.content - ): - model_response.choices[0].delta.content = "" + (model_response.choices[0].delta.content or "") - self.sent_last_thinking_block = True + if self.merge_reasoning_content_in_choices is not True: + return + delta: Final = model_response.choices[0].delta + reasoning_content: Final = getattr(delta, "reasoning_content", None) + content: Final = delta.content or "" + if reasoning_content: + self._emitted_disqualifying_content = True + closes_thinking: Final = bool(content) and not self.sent_last_thinking_block + opening: Final = "" if self.sent_first_thinking_block else "" + delta.content = opening + reasoning_content + ("" + content if closes_thinking else "") + self.sent_first_thinking_block = True + self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking + elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content: + delta.content = "" + content + self.sent_last_thinking_block = True - if hasattr(model_response.choices[0].delta, "reasoning_content"): - del model_response.choices[0].delta.reasoning_content + if hasattr(delta, "reasoning_content"): + del delta.reasoning_content def _dispatch_provider_chunk( self, 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/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 55eaa0d2ac6..ecf050d185e 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -11,7 +11,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -37,6 +37,7 @@ class AnthropicTokenCounter(BaseTokenCounter): system: JsonValue = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 60c63964492..7b8c2550838 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -139,7 +139,7 @@ def request_safeguards_evaluator( ) -async def _prepare_context_managed_request( +async def prepare_context_managed_request( *, model: str, messages: _AnthropicMessages, @@ -152,6 +152,8 @@ async def _prepare_context_managed_request( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Apply client compaction history, then optional context_management polyfill.""" from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( @@ -193,6 +195,8 @@ async def _prepare_context_managed_request( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, + drop_invalid_edits=drop_invalid_edits, ) if polyfill_result is not None: @@ -266,16 +270,14 @@ def _spec_has_non_compact_edits( return any(isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits) -def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool: +def context_management_explicitly_dropped(additional_drop_params: object) -> bool: """True when the caller opted out of context_management via ``additional_drop_params``. ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` is a LiteLLM-supported param (native on Anthropic, polyfilled elsewhere), and ``drop_params`` only exists to drop genuinely unsupported params. """ - if not isinstance(additional_drop_params, list): - return False - return "context_management" in additional_drop_params + return isinstance(additional_drop_params, list) and "context_management" in additional_drop_params def _normalize_spec_edits( @@ -291,7 +293,7 @@ def _normalize_spec_edits( if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None from litellm.llms.anthropic.pass_through.context_management.dispatcher import ( @@ -317,20 +319,23 @@ async def _run_polyfill_if_enabled( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Run the async context_management polyfill if a spec is present. Returns ``None`` when the spec is empty or ``context_management`` is listed in ``additional_drop_params`` (the explicit opt-out; ``drop_params`` - does not disable the polyfill because context_management is a supported + does not disable a valid spec because context_management is a supported param). Raises ``AnthropicContextManagementError`` so the /v1/messages - endpoint can emit an Anthropic-format 400. All other exceptions are + endpoint can emit an Anthropic-format 400, unless ``drop_invalid_edits`` + (``drop_params``) drops the invalid spec. All other exceptions are best-effort swallowed (matches v0 behavior). """ if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None try: @@ -345,13 +350,17 @@ async def _run_polyfill_if_enabled( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) except AnthropicContextManagementError: + if drop_invalid_edits: + return None # Surface validation errors so the endpoint can emit an Anthropic-format # 400. Other exception types fall into the best-effort branch below. raise except Exception as e: - verbose_logger.exception("context_management polyfill: skipping edits due to error: %s", e) + log_failure: Final = verbose_logger.debug if count_only else verbose_logger.exception + log_failure("context_management polyfill: skipping edits due to error: %s", e) # Best-effort swallow is only safe for compact-only specs, where the # caller's compaction-block-slicing safety net produces a correct # (if degraded) result. When the spec also requested non-compact @@ -709,6 +718,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: """Handle non-Anthropic models asynchronously using the adapter""" context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) litellm_router: Final[Router | None] = _router_or_proxy_fallback(kwargs.pop("litellm_router", None)) proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) @@ -716,7 +728,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: sub_call_logging_kwargs: Final = proxy_sub_call_logging_kwargs(request_kwargs) sub_call_trace_id: Final = proxy_sub_call_trace_id(request_kwargs) - polyfill_result: Final = await _prepare_context_managed_request( + polyfill_result: Final = await prepare_context_managed_request( model=model, messages=messages, tools=tools, @@ -728,6 +740,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=sub_call_trace_id, + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages @@ -841,6 +854,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: # bridge to it via ``run_async_function``. context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -866,7 +882,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) request_kwargs: Final = _REQUEST_KWARGS.validate_python(kwargs) polyfill_result = run_async_function( - _prepare_context_managed_request, + prepare_context_managed_request, model=model, messages=messages, tools=tools, @@ -878,6 +894,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=proxy_sub_call_logging_kwargs(request_kwargs), parent_trace_id=proxy_sub_call_trace_id(request_kwargs), + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 18db0f28338..52b836c789a 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,6 +24,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -39,6 +40,7 @@ from litellm.types.llms.anthropic import ( UsageDelta, UsageIteration, ) +from litellm.types.llms.openai import ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock from litellm.types.utils import AdapterCompletionStreamWrapper, ChatCompletionDeltaToolCall, Delta, Function if TYPE_CHECKING: @@ -50,6 +52,7 @@ if TYPE_CHECKING: _STREAMING_DELTA_TYPES: Final = frozenset(get_args(StreamingContentBlockDeltaType)) _SSE_EVENT: Final = TypeAdapter(Mapping[str, object]) +_THINKING_BLOCK: Final = TypeAdapter(ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) class _UsageDeltaWithIterations(UsageDelta, total=False): @@ -106,6 +109,14 @@ def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str: assert_never(delta_type) +def _first_thinking_block_index(chunk: "ModelResponseStream") -> int | None: + delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + first_block: Final = as_str_mapping(blocks[0]) if blocks else None + index: Final = first_block.get("index") if first_block is not None else None + return index if isinstance(index, int) else None + + def _optional_str_attr(obj: object, name: str) -> str | None: value: Final = _optional_attr(obj, name) return value if isinstance(value, str) else None @@ -373,13 +384,44 @@ class _CombinedChunkSplitter: finish_delta.thinking_blocks = None return [content_chunk, finish_chunk] + @staticmethod + def _with_thinking_block( + chunk: "ModelResponseStream", + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, + *, + keep_role: bool, + ) -> "ModelResponseStream": + piece: Final = copy.deepcopy(chunk) + role: Final = getattr(piece.choices[0].delta, "role", None) if keep_role else None + thinking: Final = block.get("thinking") if block["type"] == "thinking" else None + piece.choices[0].delta = Delta(role=role, reasoning_content=thinking or None, thinking_blocks=[block]) + return piece + + @staticmethod + def _split_thinking_blocks(chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": + """Anthropic streams one content block per thinking block, so a delta carrying several is split per block.""" + choices: Final = _optional_attr_sequence(chunk, "choices") + if len(choices) != 1: + return (chunk,) + delta: Final = _optional_attr(choices[0], "delta") + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + if len(blocks) < 2 or _optional_attr(delta, "content") or _optional_attr(delta, "tool_calls"): + return (chunk,) + return tuple( + _CombinedChunkSplitter._with_thinking_block( + chunk, _THINKING_BLOCK.validate_python(block), keep_role=position == 0 + ) + for position, block in enumerate(blocks) + ) + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk for combined_chunk in self._split(chunk) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) - return tuple(itertools.chain.from_iterable(self._route_tool_calls(split_chunk) for split_chunk in split_chunks)) + block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) + return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, block_chunks))) def __iter__(self) -> "Iterator[ModelResponseStream]": return self @@ -458,7 +500,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False sent_content_block_start: bool = False sent_content_block_finish: bool = False - current_content_block_type: Literal["text", "tool_use", "thinking"] = "text" + current_content_block_type: Literal["text", "tool_use", "thinking", "redacted_thinking"] = "text" + current_thinking_block_index: int | None = None sent_last_message: bool = False holding_chunk: ContentBlockDelta | None = None holding_stop_reason_chunk: MessageBlockDelta | None = None @@ -1254,13 +1297,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): Also gates re-emission after a block transition: when an upstream chunk both *triggers* a new content block (its type differs from the active block) and *carries* delta content, that content belongs to - the new block. The synthesized ``content_block_start`` only ever - carries an empty body — see - ``_translate_streaming_openai_chunk_to_anthropic_content_block``, - which returns an empty ``TextBlock``/``ToolUseBlock``/thinking block — - so the trigger chunk's delta must be re-queued or the first token of - the new block (the first non-empty text/thinking delta, or bundled - tool arguments) is silently dropped. + the new block. The synthesized ``content_block_start`` from + ``_translate_streaming_openai_chunk_to_anthropic_content_block`` + carries an empty text, tool-use or unsigned thinking body, and only a + signed thinking block's text rides on the start because its delta is + a ``signature_delta``. So the trigger chunk's delta must be re-queued + or the first token of the new block (the first non-empty text or + unsigned thinking delta, or bundled tool arguments) is silently + dropped. Delta types outside ``StreamingContentBlockDeltaType`` — the closed set the translate layer can produce — are treated as empty. The @@ -1352,11 +1396,24 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name + thinking_block_index: Final = _first_thinking_block_index(chunk) + previous_thinking_block_index: Final = self.current_thinking_block_index + if thinking_block_index is not None or block_type != self.current_content_block_type: + self.current_thinking_block_index = thinking_block_index + if block_type != self.current_content_block_type: self.current_content_block_type = block_type self.current_content_block_start = content_block_start return True + if ( + thinking_block_index is not None + and previous_thinking_block_index is not None + and thinking_block_index != previous_thinking_block_index + ): + self.current_content_block_start = content_block_start + return True + # For parallel tool calls, we'll necessarily have a new content block # if we get a function name since it signals a new tool call if block_type == "tool_use": diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index d1560c5b2ed..3aad0be1913 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -1627,7 +1627,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_streaming_openai_chunk_to_anthropic_content_block( self, choices: list[OpenAIStreamingChoice | StreamingChoices] ) -> tuple[ - Literal["text", "tool_use", "thinking"], + Literal["text", "tool_use", "thinking", "redacted_thinking"], "ContentBlockContentBlockDict", ]: from litellm._uuid import uuid @@ -1664,6 +1664,10 @@ class LiteLLMAnthropicMessagesAdapter: thinking_blocks = choice.delta.thinking_blocks or [] if len(thinking_blocks) > 0: thinking_block = thinking_blocks[0] + if thinking_block["type"] == "redacted_thinking": + return "redacted_thinking", ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", data=thinking_block.get("data") or "" + ) if thinking_block["type"] == "thinking": thinking = thinking_block.get("thinking") or "" signature = thinking_block.get("signature") or "" @@ -1672,7 +1676,7 @@ class LiteLLMAnthropicMessagesAdapter: assert isinstance(signature, str) return "thinking", ChatCompletionThinkingBlock( - type="thinking", thinking=thinking, signature=signature + type="thinking", thinking=thinking if signature else "", signature=signature ) # OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning # parsers) populate ``reasoning_content`` without ``thinking_blocks``. diff --git a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py index 27b7412292e..3e1bab09d5a 100644 --- a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py @@ -85,6 +85,7 @@ async def apply_context_management( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: "ProxySubCallLoggingKwargs | None" = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -128,6 +129,7 @@ async def apply_context_management( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) if editor_is_async else await asyncify(editor)( diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py index 09f49e9b663..5eb2863a241 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py @@ -642,6 +642,7 @@ async def apply_compact_20260112( user_api_key_auth: Optional["UserAPIKeyAuth"] = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Apply ``compact_20260112``; return a ``PolyfillResult``. @@ -682,7 +683,7 @@ async def apply_compact_20260112( # Opt-in gate: no summary model configured → no-op (but still return the # Phase A-sliced/stripped messages so compaction blocks don't leak). - summary_model: Final = _read_summary_model_setting() + summary_model: Final = None if count_only else _read_summary_model_setting() if summary_model is None: applied["error"] = "summary_model_not_configured" # Slice-only forwarding: ``augmented_system`` already carries any prior 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/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 0ed432cc394..d1f45751d3e 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -13,7 +13,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -54,6 +54,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index a8e6320bd1e..7aae2cb949d 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,14 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) +class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): + def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract + usage: Final = chunk.get("usage") + if 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: @@ -319,7 +328,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 +341,20 @@ class AzureAIStudioConfig(OpenAIConfig): api_key=api_key, json_mode=json_mode, ) + 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/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index aede8aedc60..611051bb2dd 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -6,16 +6,38 @@ import litellm if TYPE_CHECKING: import httpx +from litellm.types.llms.openai import ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock from litellm.types.utils import ( Choices, Delta, GenericStreamingChunk, + Message, ModelResponse, ModelResponseStream, StreamingChoices, ) +def _indexed_thinking_block( + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, index: int +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: + if block["type"] == "redacted_thinking": + return {**block, "index": index} + return {**block, "index": index} + + +def _streaming_delta(message: Message) -> Delta: + thinking_blocks: Final = message.thinking_blocks if hasattr(message, "thinking_blocks") else None + return Delta( + **message.model_dump(exclude={"thinking_blocks"}), + thinking_blocks=( + None + if thinking_blocks is None + else [_indexed_thinking_block(block, index) for index, block in enumerate(thinking_blocks)] + ), + ) + + def convert_model_response_to_streaming( model_response: ModelResponse, ) -> ModelResponseStream: @@ -40,9 +62,7 @@ def convert_model_response_to_streaming( streaming_choices.append( StreamingChoices( index=choice.index, - delta=Delta( - **cast(Choices, choice).message.model_dump(), - ), + delta=_streaming_delta(cast(Choices, choice).message), finish_reason=choice.finish_reason, ) ) diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 70ec6572675..5ed73e13188 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -16,7 +16,13 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + Message, + ProviderSpecificModelInfo, + TokenCountMessagesFormat, + TokenCountResponse, +) class BaseTokenCounter(ABC): @@ -32,6 +38,7 @@ class BaseTokenCounter(ABC): system: Any | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: pass 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..bbec66c6765 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -390,8 +390,24 @@ def systemone_response_to_ir(response: DecisionsResponse, request: DecisionsIRRe ) +_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_LABELS: Final[TypeAdapter[tuple[str, ...]]] = TypeAdapter(tuple[str, ...]) + + +def _normalize_answer(value: object) -> Mapping[str, object]: + answer: Final = _MAPPING.validate_python(value) + legend: Final = answer.get("legend") + if answer.get("type") != "score" or not isinstance(legend, list): + return answer + labels: Final = _LABELS.validate_python(legend) + return {**answer, "legend": {str(index): label for index, label in enumerate(labels)}} + + def parse_systemone_response(payload: object, request: DecisionsIRRequest) -> DecisionsIRResponse: - return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(payload), request) + response: Final = _MAPPING.validate_python(payload) + answers: Final = _MAPPING.validate_python(response.get("answers")) + normalized: Final = {**response, "answers": {name: _normalize_answer(answer) for name, answer in answers.items()}} + return systemone_response_to_ir(_SYSTEMONE_RESPONSE_ADAPTER.validate_python(normalized), request) def _systemone_answer(answer: DecisionsIRAnswer) -> DecisionAnswer | None: @@ -495,6 +511,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/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py index c3264911d46..b4514fd1da2 100644 --- a/litellm/llms/base_llm/realtime/transcription_protocol.py +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -102,6 +102,16 @@ def json_integer( return value +def uses_server_vad( + update: TranscriptionSessionUpdate, backend: str, drop_params: bool, error: ProtocolErrorType +) -> bool: + if update.turn_detection_disabled: + return False + if update.turn_detection_type in (None, "server_vad") or drop_params: + return True + raise error(f"{backend} supports server_vad turn detection or null. Set drop_params to use server_vad instead") + + def new_event_id() -> str: return f"event_{uuid.uuid4().hex}" diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 35d08d4b2c2..5a64e3fb78e 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -10,6 +10,7 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.llms.base_llm import base_model_iterator from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -21,7 +22,7 @@ from litellm.utils import CustomStreamWrapper from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, pop_aws_auth_params, run_aws_signing from ..common_utils import BedrockError, error_response_text, get_all_bedrock_regions -from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call +from .invoke_handler import AWSEventStreamDecoder, make_call def make_sync_call( @@ -36,7 +37,8 @@ def make_sync_call( fake_stream: bool = False, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> tuple[MockResponseIterator | Iterator[object], httpx.Headers]: + optional_params: Mapping[str, object] | None = None, +) -> tuple[base_model_iterator.MockResponseIterator | Iterator[object], httpx.Headers]: if client is None: client = get_httpx_client() # Create a new client if none provided @@ -64,14 +66,15 @@ def make_sync_call( model_response=litellm.ModelResponse(), stream=True, logging_obj=logging_obj, - optional_params={}, + optional_params=dict(optional_params or {}), api_key="", data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | Iterator[object] = MockResponseIterator( - model_response=model_response, json_mode=json_mode + completion_stream: base_model_iterator.MockResponseIterator | Iterator[object] = ( + base_model_iterator.MockResponseIterator(model_response=model_response) ) else: decoder: Final = AWSEventStreamDecoder(model=model, json_mode=json_mode) @@ -158,6 +161,7 @@ class BedrockConverseLLM(BaseAWSLLM): json_mode=json_mode, stream_chunk_size=stream_chunk_size, timeout=timeout, + optional_params=optional_params, ) streaming_response: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -290,7 +294,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream_chunk_size: Final = stored_control_options(litellm_params).stream_chunk_size if stream is True else None unencoded_model_id: Final = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) - json_mode: Final = optional_params.get("json_mode", False) + json_mode: Final = optional_params.get("json_mode") is True if unencoded_model_id is not None: modelId = self.encode_model_id(model_id=unencoded_model_id) else: @@ -463,6 +467,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream=fake_stream, stream_chunk_size=stream_chunk_size, timeout=timeout, + optional_params=optional_params, ) streaming_response: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 6ad400661be..7553dacd35d 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1,7 +1,7 @@ import json import types from collections.abc import AsyncIterator, Iterator, Mapping -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter, ValidationError @@ -20,6 +20,7 @@ from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import JsonModeRequest +from litellm.llms.base_llm import base_model_iterator from litellm.llms.base_llm.base_utils import ResponseFormatSchema from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -32,12 +33,8 @@ from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, - ChatCompletionUsageBlock, ) from litellm.types.utils import ( - ChatCompletionMessageToolCall, - Choices, Delta, ModelResponse, ModelResponseStream, @@ -228,7 +225,10 @@ async def make_call( bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> "tuple[MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict], httpx.Headers]": + optional_params: Mapping[str, object] | None = None, +) -> tuple[ + base_model_iterator.MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict], httpx.Headers +]: try: if client is None: client = get_async_httpx_client( @@ -264,15 +264,16 @@ async def make_call( model_response=litellm.ModelResponse(), stream=True, logging_obj=logging_obj, - optional_params={}, + optional_params=dict(optional_params or {}), api_key="", data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] = ( - MockResponseIterator(model_response=model_response, json_mode=json_mode) - ) + completion_stream: ( + base_model_iterator.MockResponseIterator | AsyncIterator[GChunk | ModelResponseStream | dict] + ) = base_model_iterator.MockResponseIterator(model_response=model_response) else: decoder: Final = _stream_decoder( bedrock_invoke_provider, @@ -324,7 +325,7 @@ def make_sync_call( bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, timeout: float | httpx.Timeout | None = None, -) -> "tuple[MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict], httpx.Headers]": +) -> tuple[base_model_iterator.MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict], httpx.Headers]: try: if client is None: client = get_httpx_client( @@ -364,10 +365,11 @@ def make_sync_call( data=data, messages=messages, encoding=litellm.encoding, + json_mode=json_mode, ) - completion_stream: MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] = ( - MockResponseIterator(model_response=model_response, json_mode=json_mode) - ) + completion_stream: ( + base_model_iterator.MockResponseIterator | Iterator[GChunk | ModelResponseStream | dict] + ) = base_model_iterator.MockResponseIterator(model_response=model_response) else: decoder: Final = _stream_decoder( bedrock_invoke_provider, @@ -991,98 +993,3 @@ class AmazonOpenAICompatibleStreamDecoder(AWSEventStreamDecoder): def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: return self.openai_model_response_iterator.chunk_parser(chunk=chunk_data) - - -class MockResponseIterator: # for returning ai21 streaming responses - def __init__(self, model_response, json_mode: bool | None = False): - self.model_response = model_response - self.json_mode = json_mode - self.is_done = False - - # Sync iterator - def __iter__(self): - return self - - def _handle_json_mode_chunk( - self, text: str, tool_calls: list[ChatCompletionToolCallChunk] | None - ) -> tuple[str, ChatCompletionToolCallChunk | None]: - """ - If JSON mode is enabled, convert the tool call to a message. - - Bedrock returns the JSON schema as part of the tool call - OpenAI returns the JSON schema as part of the content, this handles placing it in the content - - Args: - text: str - tool_use: Optional[ChatCompletionToolCallChunk] - Returns: - Tuple[str, Optional[ChatCompletionToolCallChunk]] - - text: The text to use in the content - tool_use: The ChatCompletionToolCallChunk to use in the chunk response - """ - tool_use: ChatCompletionToolCallChunk | None = None - if self.json_mode is True and tool_calls is not None: - message: Final = litellm.AnthropicConfig().convert_tool_response_to_message(tool_calls=tool_calls) - if message is not None: - text = message.content or "" - tool_use = None - elif tool_calls is not None and len(tool_calls) > 0: - tool_use = tool_calls[0] - return text, tool_use - - def _chunk_parser(self, chunk_data: ModelResponse) -> GChunk: - try: - chunk_usage: Final[Usage] = getattr(chunk_data, "usage") - text = chunk_data.choices[0].message.content or "" - tool_use = None - _model_response_tool_call: Final = cast( - list[ChatCompletionMessageToolCall] | None, - cast(Choices, chunk_data.choices[0]).message.tool_calls, - ) - if self.json_mode is True: - text, tool_use = self._handle_json_mode_chunk( - text=text, - tool_calls=chunk_data.choices[0].message.tool_calls, - ) - elif _model_response_tool_call is not None: - tool_use = ChatCompletionToolCallChunk( - id=_model_response_tool_call[0].id, - type="function", - function=ChatCompletionToolCallFunctionChunk( - name=_model_response_tool_call[0].function.name, - arguments=_model_response_tool_call[0].function.arguments, - ), - index=0, - ) - processed_chunk: Final = GChunk( - text=text, - tool_use=tool_use, - is_finished=True, - finish_reason=map_finish_reason(finish_reason=chunk_data.choices[0].finish_reason or ""), - usage=ChatCompletionUsageBlock( - prompt_tokens=chunk_usage.prompt_tokens, - completion_tokens=chunk_usage.completion_tokens, - total_tokens=chunk_usage.total_tokens, - ), - index=0, - ) - return processed_chunk - except Exception as e: - raise ValueError(f"Failed to decode chunk: {chunk_data}. Error: {e}") - - def __next__(self): - if self.is_done: - raise StopIteration - self.is_done = True - return self._chunk_parser(self.model_response) - - # Async iterator - def __aiter__(self): - return self - - async def __anext__(self): - if self.is_done: - raise StopAsyncIteration - self.is_done = True - return self._chunk_parser(self.model_response) 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/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 7ff2fddbb57..97f28b9fbc5 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,13 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + LiteLLMPydanticObjectBase, + LlmProviders, + TokenCountMessagesFormat, + TokenCountResponse, +) RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -109,6 +115,7 @@ class BedrockTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: if not messages: return None 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/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index c0a382b136c..3e7c2f7eee9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -8,7 +8,7 @@ import asyncio import contextlib import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, NoReturn, Protocol, runtime_checkable @@ -21,7 +21,7 @@ from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, - BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, @@ -37,9 +37,8 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities -_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -177,25 +176,35 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: return -def _pending_session_update(scope: Mapping[str, object]) -> str | None: - """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: - committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) - raise BedrockError( - status_code=400, - message=( - "Bedrock realtime session already committed to a provider stream; it cannot be replayed" - + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") - ), - ) - pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) +def _is_session_committed(scope: Mapping[str, object]) -> bool: + return scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True + + +def raise_if_session_committed(scope: Mapping[str, object] | None) -> None: + """A fallback attempt on a websocket whose Bedrock stream already took input would silently start over.""" + if scope is None or not _is_session_committed(scope): + return + committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) + raise BedrockError( + status_code=400, + message=( + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" + + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + ), + ) + + +def _pending_opening_message(scope: Mapping[str, object]) -> str | None: + """A fallback attempt on the same websocket replays the opening message Bedrock never accepted a stream for.""" + raise_if_session_committed(scope) + pending: Final = scope.get(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY) return pending if isinstance(pending, str) else None def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: error: Final = _as_bedrock_error(failure) verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error))) - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + if _is_session_committed(scope): scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error)) raise error from failure @@ -207,26 +216,36 @@ def _parse_client_message(message: str) -> Mapping[str, JsonValue]: return _EMPTY_JSON_OBJECT -async def _ack_session_update( +async def _send_opening_message( client_ws: RealtimeClientWebSocket, bedrock_stream: BedrockBidirectionalStream, + message: str, + bedrock_messages: Sequence[str], + send_to_bedrock: Callable[[str], Awaitable[None]], +) -> None: + """Keep the opening message for a fallback to replay until Bedrock accepts the stream it opened.""" + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = message # rebind-ok: scope outlives attempts + for bedrock_message in bedrock_messages: + await send_to_bedrock(bedrock_message) + await bedrock_stream.await_output() + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging | None, parsed_client_message: Mapping[str, JsonValue], ) -> bool: - """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" - await bedrock_stream.await_output() - client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) - client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + """Ack the client's session.update; False means the client is gone.""" if logging_obj is None: return True - requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) + modalities: Final = requested_modalities(_json_dict(parsed_client_message.get("session"))) try: await client_ws.send_text( - json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + json.dumps(transformation_config.session_updated_event(model, logging_obj, modalities)) ) except Exception as e: # noqa: BLE001 # any send failure means the client is gone verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) @@ -314,7 +333,7 @@ class BedrockRealtime(BaseAWSLLM): except ImportError as e: raise _sdk_import_error(self._sdk_version_lookup(), e) from e - pending_session_update: Final = _pending_session_update(websocket.scope) + pending_opening_message: Final = _pending_opening_message(websocket.scope) # Get AWS region if aws_region_name is None: @@ -376,7 +395,7 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params ) finally: await _close_bedrock_client(bedrock_client) @@ -387,7 +406,7 @@ class BedrockRealtime(BaseAWSLLM): open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]], model: str, logging_obj: LiteLLMLogging, - pending_session_update: str | None, + pending_opening_message: str | None, drop_params: bool = False, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -396,7 +415,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") - if pending_session_update is None: + if pending_opening_message is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") @@ -418,7 +437,7 @@ class BedrockRealtime(BaseAWSLLM): model, session_state, logging_obj, - initial_message=pending_session_update, + initial_message=pending_opening_message, ) logged_events: Final = ( @@ -531,18 +550,22 @@ class BedrockRealtime(BaseAWSLLM): if not await _reject_session_update(client_ws, parsed_client_message, rejection): break continue - client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message( + client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock + ) + else: + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) if is_session_update and not await _ack_session_update( - client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + client_ws, transformation_config, model, logging_obj, parsed_client_message ): break finally: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index e8bc0482a28..9cb38822881 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -123,26 +123,81 @@ def _integer_setting(value: JsonValue) -> int | None: return None if number is None else int(number) -def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: +def requested_modalities(session_config: Mapping[str, JsonValue]) -> tuple[str, ...] | None: + modalities: Final = session_config.get("modalities") + if not isinstance(modalities, list): + return None + names: Final = tuple(modality for modality in modalities if isinstance(modality, str)) + return names if len(names) == len(modalities) else None + + +_OPENAI_VOICES: Final = frozenset( + ("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar") +) + + +def _unplayable_voice(session_config: Mapping[str, JsonValue]) -> str | None: + voice: Final = session_config.get("voice") + if not isinstance(voice, str): + return None + return voice if not voice.strip() or voice.lower() in _OPENAI_VOICES else None + + +def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + input_audio_format: Final = session_config.get("input_audio_format") + output_audio_format: Final = session_config.get("output_audio_format") + modalities: Final = session_config.get("modalities") temperature: Final = session_config.get("temperature") max_output_tokens: Final = session_config.get("max_response_output_tokens") + input_rate_hertz: Final = session_config.get("input_sample_rate_hertz") output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + tools: Final = session_config.get("tools") malformed: Final = ( - ("temperature", temperature is not None and _number_setting(temperature) is None), + ("instructions", "a string", instructions is not None and not isinstance(instructions, str)), + ("voice", "a string", voice is not None and not isinstance(voice, str)), + ("tools", "a list", tools is not None and not isinstance(tools, list)), + ( + "input_audio_format", + "a string", + input_audio_format is not None and not isinstance(input_audio_format, str), + ), + ( + "output_audio_format", + "a string", + output_audio_format is not None and not isinstance(output_audio_format, str), + ), + ( + "modalities", + "a list of strings", + modalities is not None and requested_modalities(session_config) is None, + ), + ("temperature", "a number", temperature is not None and _number_setting(temperature) is None), ( "max_response_output_tokens", + "a number or 'inf'", max_output_tokens is not None and max_output_tokens != "inf" and _integer_setting(max_output_tokens) is None, ), - ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ( + "input_sample_rate_hertz", + "a number", + input_rate_hertz is not None and _integer_setting(input_rate_hertz) is None, + ), + ( + "output_sample_rate_hertz", + "a number", + output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None, + ), ) - return tuple(name for name, is_malformed in malformed if is_malformed) + return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed) @dataclasses.dataclass(frozen=True, slots=True) class BedrockSessionUpdateRejection: - code: Literal["invalid_type", "session_update_not_applied"] + code: Literal["invalid_type", "invalid_value", "session_update_not_applied"] message: str param: str @@ -312,7 +367,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return 8000 # G.711 typically uses 8kHz return 24000 if is_output else 16000 - def transform_session_update_event(self, json_message: dict) -> list[str]: + def transform_session_update_event(self, json_message: Mapping[str, JsonValue]) -> list[str]: """ Transform session.update event to Bedrock session configuration. @@ -325,7 +380,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling session.update") messages: Final[list[str]] = [] - session_config: Final = json_message.get("session", {}) + requested_session: Final = json_message.get("session") + session_config: Final[Mapping[str, JsonValue]] = ( + requested_session if isinstance(requested_session, dict) else {} + ) locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) # Settings Bedrock locks at prompt start may only change before it @@ -336,13 +394,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided - if "input_audio_format" in session_config: - input_format: Final = session_config["input_audio_format"] + input_format: Final = session_config.get("input_audio_format") + if isinstance(input_format, str): self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) # Allow direct override of the input sample rate if provided (custom extension) - if "input_sample_rate_hertz" in session_config: - self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + requested_input_rate: Final = _integer_setting(session_config.get("input_sample_rate_hertz")) + if requested_input_rate is not None: + self.input_sample_rate_hertz = requested_input_rate if self.prompt_started: return messages @@ -352,7 +411,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages.append(json.dumps(session_start)) # Send prompt start - tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + tools: Final = session_config.get("tools") prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, @@ -370,7 +429,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolUseOutputConfiguration": {"mediaType": "application/json"}, "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, } - if tools + if isinstance(tools, list) and tools else {} ), } @@ -445,12 +504,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if isinstance(requested_tools, list) else base.tools_json ) - system_prompt: Final = base.system_prompt if instructions is None else str(instructions) - voice_id: Final = base.voice_id if voice is None else str(voice) + system_prompt: Final = instructions if isinstance(instructions, str) else base.system_prompt + voice_id: Final = ( + voice if isinstance(voice, str) and _unplayable_voice(session_config) is None else base.voice_id + ) format_sample_rate: Final = ( - base.output_sample_rate_hertz - if output_audio_format is None - else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + self._map_audio_format_to_sample_rate(output_audio_format, is_output=True) + if isinstance(output_audio_format, str) + else base.output_sample_rate_hertz ) requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) output_sample_rate_hertz: Final = ( @@ -476,16 +537,26 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): session_config: Final = json_message.get("session") if not isinstance(session_config, Mapping): return None - malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + malformed: Final = () if self.drop_params else _malformed_settings(session_config) if malformed: + name, expected = malformed[0] return BedrockSessionUpdateRejection( code="invalid_type", message=( - f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " - "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. " "Set drop_params to drop it instead" ), - param=f"session.{malformed[0]}", + param=f"session.{name}", + ) + unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config) + if unplayable_voice is not None: + return BedrockSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. " + "Set drop_params to drop it instead" + ), + param="session.voice", ) if not self.prompt_started: return None @@ -785,7 +856,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, object]] = json.loads(message) + json_message: Final[dict[str, JsonValue]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -816,11 +887,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSession: session: Final = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=modalities if modalities is not None else ["text", "audio"], + modalities=list(modalities) if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model @@ -842,7 +913,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSessionEvents: """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" return OpenAIRealtimeStreamSessionEvents( 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 610321fe1f6..38e076d650a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7127,7 +7127,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, @@ -7586,7 +7587,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, @@ -7781,7 +7783,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, @@ -7961,7 +7964,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="", @@ -8124,7 +8128,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="", @@ -8264,7 +8269,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="", @@ -8436,7 +8442,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="", @@ -8622,7 +8629,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="", @@ -8754,7 +8762,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="", @@ -8835,7 +8844,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="", @@ -9019,7 +9029,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="", @@ -9122,6 +9133,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 @@ -9197,9 +9209,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 @@ -9293,6 +9307,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: @@ -9369,9 +9384,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: @@ -9460,6 +9477,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: @@ -9536,9 +9554,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: @@ -9629,6 +9649,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: @@ -9705,9 +9726,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: @@ -9803,6 +9826,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: @@ -9881,9 +9905,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: @@ -9972,6 +9998,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: @@ -10047,9 +10074,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: @@ -10167,8 +10196,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: @@ -10410,8 +10441,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: @@ -10541,8 +10574,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: @@ -10663,8 +10698,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: @@ -10805,8 +10842,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: @@ -10950,8 +10989,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: @@ -11068,8 +11109,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: @@ -11210,8 +11253,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) @@ -11351,8 +11396,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) @@ -11482,8 +11529,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) @@ -11619,8 +11668,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) @@ -11762,8 +11813,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/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 3d04954dd41..4006c89954d 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -497,12 +497,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, + count_context_edits, count_context_error_response, gemini_count_contents, inference_request_view, @@ -511,7 +513,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): vertex_gemini_count_context, ) - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 7c04ebf311f..f4e2941670a 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -161,6 +161,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" return True + def _rejects_vad_tuning_with_detection_off(self) -> bool: + """Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``.""" + return True + @staticmethod def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: if not isinstance(details, dict): @@ -279,10 +283,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): OpenAI ``semantic_vad`` has no Gemini Live equivalent — return an empty dict so callers omit ``realtimeInputConfig`` (mapping it with - ``disabled: true`` breaks native-audio sessions). + ``disabled: true`` breaks native-audio sessions). An explicit + ``create_response: false`` maps to detection off alone where the backend + rejects tuning next to ``disabled: true``. """ if isinstance(value, dict) and value.get("type") == "semantic_vad" and "create_response" not in value: return AutomaticActivityDetection() + auto_response_off: Final = "create_response" in value and value["create_response"] is False + if auto_response_off and self._rejects_vad_tuning_with_detection_off(): + return AutomaticActivityDetection(disabled=True) automatic_activity_dection: Final = AutomaticActivityDetection() if "create_response" in value and isinstance(value["create_response"], bool): diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 46231e5e980..892d7ff66de 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -26,6 +26,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.secret_managers.main import get_secret_str @@ -230,15 +231,7 @@ def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate: return 16_000 if rate == 16_000 else 24_000 -def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: - if update.turn_detection_disabled: - return "PUSH_TO_TALK" - if update.turn_detection_type not in (None, "server_vad"): - raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") - return "ENDPOINTING" - - -def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: +def parse_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> MuseSessionConfig: update: Final = parse_transcription_session_update(payload, MuseProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") @@ -253,7 +246,7 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig raise MuseProtocolError("realtime session model cannot be changed") return MuseSessionConfig( model=normalized_model, - mode=_parse_mode(update), + mode="ENDPOINTING" if uses_server_vad(update, "Muse Voice", drop_params, MuseProtocolError) else "PUSH_TO_TALK", sample_rate=_parse_sample_rate(update), language_bias=() if update.language is None else (normalize_language(update.language),), ) @@ -441,9 +434,11 @@ class MetaRealtimeConfig(BaseRealtimeConfig): *, monotonic: Callable[[], float] = time.monotonic, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + drop_params: bool = False, ) -> None: self._monotonic: Final = monotonic self._sleep: Final = sleep + self._drop_params: Final = drop_params self._transformer: Final = MuseEventTransformer() self._access_token: str | None = None self._config: MuseSessionConfig | None = None @@ -553,7 +548,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): access_token: Final = self._access_token if access_token is None: raise MuseProtocolError("Meta API key was not validated before the session was configured") - config: Final = parse_session_update(message, model) + config: Final = parse_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config) return (json.dumps(config.handshake(access_token), separators=(",", ":")),) diff --git a/litellm/llms/nace/decisions/transformation.py b/litellm/llms/nace/decisions/transformation.py new file mode 100644 index 00000000000..c96177bb7d5 --- /dev/null +++ b/litellm/llms/nace/decisions/transformation.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class NaceDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("NACE_API_KEY",) + api_base_env = ("NACE_API_BASE", "NACE_BASE_URL") + + def get_default_api_base(self) -> str | None: + return "https://console.nace.ai" 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 27346ee936e..73f40f5e0e4 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -11,13 +11,14 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer import openai -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai._base_client import make_request_options from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted +from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, TypeAdapter from typing_extensions import overload @@ -30,9 +31,8 @@ from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, MockResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException -from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator from litellm.types.utils import ( EmbeddingResponse, ImageResponse, @@ -53,16 +53,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() @@ -352,7 +355,14 @@ _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) -def _embedding_request_without_sdk_defaults( +def _raw_chat_response( + result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk], +) -> LegacyAPIResponse[ChatCompletion]: + assert isinstance(result, LegacyAPIResponse) + return result + + +def _request_without_sdk_transform( data: Mapping[str, object], timeout: float | httpx.Timeout ) -> tuple[Mapping[str, object], RequestOptions]: body: Final = {k: v for k, v in data.items() if k not in _SDK_OPTION_KEYS} @@ -426,7 +436,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, @@ -448,7 +458,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, @@ -515,7 +525,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ start_time: Final = time.time() try: - raw_response = await openai_aclient.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + await openai_aclient.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=AsyncStream[ChatCompletionChunk], + ) + ) end_time = time.time() if hasattr(raw_response, "headers"): @@ -556,7 +576,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ raw_response = None try: - raw_response = openai_client.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + openai_client.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=Stream[ChatCompletionChunk], + ) + ) if hasattr(raw_response, "headers"): headers = dict(raw_response.headers) @@ -1255,7 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = await openai_aclient.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) @@ -1272,7 +1302,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = openai_client.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) @@ -1456,10 +1486,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, @@ -1521,6 +1553,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: @@ -1542,10 +1575,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, @@ -1623,6 +1658,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( @@ -1639,6 +1675,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( @@ -1649,6 +1686,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) @@ -1686,6 +1724,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, @@ -1697,6 +1736,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ), ) @@ -1743,26 +1783,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, @@ -1782,6 +1823,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, @@ -1791,6 +1833,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1824,6 +1867,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, @@ -1833,6 +1877,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1889,6 +1934,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, @@ -1898,6 +1944,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1952,6 +1999,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, @@ -1961,6 +2009,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1998,6 +2047,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, @@ -2007,6 +2057,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2047,6 +2098,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, @@ -2056,6 +2108,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2101,26 +2154,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, @@ -2140,6 +2194,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, @@ -2149,6 +2204,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2184,6 +2240,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, @@ -2193,6 +2250,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2227,6 +2285,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, @@ -2236,6 +2295,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2276,6 +2336,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, @@ -2285,6 +2346,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2313,22 +2375,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, @@ -2338,22 +2397,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 ### @@ -2369,6 +2425,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, @@ -2376,6 +2434,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) request_params: Final = { @@ -2403,6 +2462,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]]: ... @@ -2416,6 +2477,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_assistants: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[Assistant]: ... @@ -2434,7 +2497,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, @@ -2442,6 +2507,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( @@ -2450,6 +2516,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2477,6 +2544,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, @@ -2484,6 +2553,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2501,6 +2571,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( @@ -2509,6 +2581,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, create_assistant_data=create_assistant_data, ) @@ -2518,6 +2591,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2534,6 +2608,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, @@ -2541,6 +2617,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2558,6 +2635,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( @@ -2566,6 +2645,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, assistant_id=assistant_id, ) @@ -2575,6 +2655,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2593,6 +2674,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, @@ -2600,6 +2683,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2630,6 +2714,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]: ... @@ -2645,6 +2731,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, a_add_message: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIMessage: ... @@ -2661,6 +2749,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( @@ -2671,6 +2761,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( @@ -2679,6 +2770,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2704,6 +2796,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, @@ -2711,6 +2805,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2731,6 +2826,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]]: ... @@ -2745,6 +2842,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_messages: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[OpenAIMessage]: ... @@ -2760,6 +2859,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( @@ -2769,6 +2870,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( @@ -2777,6 +2879,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2796,6 +2899,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, @@ -2803,6 +2908,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2835,6 +2941,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]: ... @@ -2850,6 +2958,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client: OpenAI | None, acreate_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2866,6 +2976,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client=None, acreate_thread=None, + *, + litellm_params: Mapping[str, object] | None = None, ): """ Here's an example: @@ -2885,6 +2997,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, messages=messages, ) @@ -2894,6 +3007,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2921,6 +3035,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, @@ -2928,6 +3044,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2953,6 +3070,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]: ... @@ -2967,6 +3086,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2982,6 +3103,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( @@ -2991,6 +3114,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( @@ -2999,6 +3123,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3032,6 +3157,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, @@ -3039,6 +3166,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3143,6 +3271,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[True], event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Run]: ... @@ -3165,6 +3295,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[False] | None, event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Run: ... @@ -3188,6 +3320,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: @@ -3197,6 +3331,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) return self.async_run_thread_stream( @@ -3224,6 +3359,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( @@ -3232,6 +3368,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 7c0697626ca..c7c997bc961 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -27,21 +27,27 @@ 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 LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests 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): @@ -95,6 +101,7 @@ class OpenAITokenCounter(BaseTokenCounter): system: object = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -125,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/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 62a84ad3f15..2c74ef19e7d 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -14,7 +14,7 @@ from typing_extensions import ReadOnly import litellm from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.databricks.streaming_utils import ModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index 83fcb43c9b8..e07b5627256 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage from litellm.types.utils import ModelResponse @@ -119,6 +120,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): choice["message"] = message returned_response: Final = ModelResponse(**response_json) + set_provider_response_headers_in_hidden_params(returned_response, response.headers) if custom_llm_provider is not None: returned_response.model = custom_llm_provider + "/" + (returned_response.model or "") diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index bf2a7d7da15..6a480904a3e 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict import httpx from typing_extensions import ReadOnly +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( anthropic_process_openai_file_message, convert_to_anthropic_tool_result, @@ -184,10 +185,11 @@ def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]: # mutable- blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None) if not isinstance(blocks, list): return [] + mappings: Final = tuple(filter(None, map(as_str_mapping, blocks))) return [ - dict(block) - for block in blocks - if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking") + {key: value for key, value in block.items() if key != "index"} + for block in mappings + if block.get("signature") or block.get("type") == "redacted_thinking" ] diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 1d4a974fc54..5015678c467 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -13,7 +13,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.realtime.transcription_protocol import ( RealtimeTranscriptionProtocolError, TranscriptionAudioFormat, - TranscriptionSessionUpdate, completed_event, decode_pcm16_append, delta_event, @@ -23,6 +22,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend from litellm.llms.vertex_ai.audio_transcription.transformation import ( @@ -125,7 +125,7 @@ def default_session_config(model: str) -> ChirpSessionConfig: ) -def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: +def parse_chirp_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> ChirpSessionConfig: update: Final = parse_transcription_session_update(payload, ChirpProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") @@ -141,7 +141,7 @@ def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessio model=model, language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language), sample_rate=_parse_sample_rate(update.audio_format), - server_vad=_parse_server_vad(update), + server_vad=uses_server_vad(update, "Speech-to-Text streaming", drop_params, ChirpProtocolError), ) @@ -161,14 +161,6 @@ def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int: return rate -def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool: - if update.turn_detection_disabled: - return False - if update.turn_detection_type not in (None, "server_vad"): - raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null") - return True - - def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: return transcription_session_created_event(config.openai_session(session_id)) @@ -331,11 +323,13 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): project: str, location: str | None, backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, + drop_params: bool = False, ) -> None: self._resolve_access_token: Final = resolve_access_token self._project: Final = validate_vertex_transcription_project_id(project) self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) self._backend_factory: Final = backend_factory + self._drop_params: Final = drop_params self._transformer: Final = ChirpEventTransformer() self._config: ChirpSessionConfig | None = None self._session_id: str | None = None @@ -423,7 +417,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): if self._config is not None: verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured") return () - config: Final = parse_chirp_session_update(message, model) + config: Final = parse_chirp_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}") return (config.configure_command(),) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 581e2cb15f7..27ca4acf05d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1274,6 +1274,7 @@ class VertexAITokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy @@ -1352,6 +1353,7 @@ class VertexAITokenCounter(BaseTokenCounter): else: from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, + count_context_edits, count_context_error_response, gemini_count_contents, inference_request_view, @@ -1361,7 +1363,14 @@ class VertexAITokenCounter(BaseTokenCounter): ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index e094f4c5c0c..871532b7af5 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -10,11 +10,10 @@ from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse if TYPE_CHECKING: from litellm.exceptions import APIError - from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -26,6 +25,7 @@ class VertexGeminiCountContext: @dataclass(frozen=True, slots=True) class VertexGeminiCountContextError: error_message: str + status_code: int = 400 _SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) @@ -39,6 +39,11 @@ _VERTEX_TOOL_KEYS: Final = MappingProxyType( ) +_INVALID_SYSTEM: Final = VertexGeminiCountContextError( + error_message="Gemini system instructions only support text blocks." +) + + def _system_parts( system: object | None, ) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: @@ -49,7 +54,7 @@ def _system_parts( try: blocks: Final = _SYSTEM_BLOCKS.validate_python(system) except ValidationError: - return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return _INVALID_SYSTEM return tuple({"text": block["text"]} for block in blocks) @@ -157,7 +162,7 @@ def count_context_error_response( tokenizer_type="gemini", error=True, error_message=error.error_message, - status_code=400, + status_code=error.status_code, ) @@ -179,36 +184,120 @@ class InferenceRequestView: system: object | None -def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": +@dataclass(frozen=True, slots=True) +class CountContextEdits: + spec: JsonValue + model: str + tools: tuple[Mapping[str, object], ...] | None + drop_invalid_edits: bool = False + + +_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( + Mapping[str, object] | tuple[Mapping[str, object], ...] +) + + +def count_context_edits( + context_management: CountContextManagement | None, + tools: Sequence[Mapping[str, object]] | None, + litellm_params: Mapping[str, object], + fallback_model: str, +) -> CountContextEdits: + model: Final = litellm_params.get("model") + return CountContextEdits( + spec=None if context_management is None else context_management.spec, + model=model if isinstance(model, str) else fallback_model, + tools=None if tools is None else tuple(tools), + drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, + ) + + +def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: + try: + return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) + except ValidationError: + return None + + +def _typed_system( + system: object | None, +) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError: + if system is None or isinstance(system, str): + return system + try: + return _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return _INVALID_SYSTEM + + +def _view_without_edits( + messages: Sequence[Mapping[str, object]], + system: object | None, + typed_system: str | Sequence[Mapping[str, object]] | None, +) -> InferenceRequestView: from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( apply_client_compaction_block_history, ) - chat_messages: Final = [dict(message) for message in messages] - if system is None or isinstance(system, str): - return apply_client_compaction_block_history(messages=chat_messages, system=system) - try: - blocks: Final = _SYSTEM_BLOCKS.validate_python(system) - except ValidationError: - return None - return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + history: Final = apply_client_compaction_block_history( + messages=[dict(message) for message in messages], + system=typed_system + if typed_system is None or isinstance(typed_system, str) + else [dict(block) for block in typed_system], + ) + if history is None: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestView(messages=history.messages, system=history.system) -def inference_request_view( - messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat -) -> InferenceRequestView: - """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" +async def _anthropic_inference_view( + messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits +) -> InferenceRequestView | VertexGeminiCountContextError: + from litellm.llms.anthropic.pass_through.adapters.handler import prepare_context_managed_request + from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + typed_system: Final = _typed_system(system) + if isinstance(typed_system, VertexGeminiCountContextError): + return typed_system + normalized: Final = normalize_native_messages(messages) + spec: Final = _typed_spec(edits.spec) + try: + prepared: Final = await prepare_context_managed_request( + model=edits.model, + messages=normalized, + tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], + system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, + context_management_spec=( + None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] + ), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + count_only=True, + drop_invalid_edits=edits.drop_invalid_edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return _view_without_edits(normalized, system, typed_system) + return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code) + if prepared is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=prepared.messages, system=prepared.system) + + +async def inference_request_view( + messages: Sequence[Mapping[str, object]], + system: object | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextEdits, +) -> InferenceRequestView | VertexGeminiCountContextError: + """Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini.""" match messages_format: case "openai": return InferenceRequestView(messages=messages, system=system) case "anthropic": - normalized: Final = normalize_native_messages(messages) - compacted: Final = _client_compaction(normalized, system) - if compacted is None: - return InferenceRequestView(messages=normalized, system=system) - return InferenceRequestView(messages=compacted.messages, system=compacted.system) + return await _anthropic_inference_view(messages, system, edits) case _: assert_never(messages_format) diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index c14dd535145..55f1cf58cbe 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -47,6 +47,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): def _include_function_response_id(self) -> bool: return False + def _rejects_vad_tuning_with_detection_off(self) -> bool: + return False + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ @@ -248,9 +251,12 @@ def vertex_realtime_config( resolve_access_token: Callable[[], Awaitable[str]], project: str, location: str | None, + drop_params: bool, ) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: if is_vertex_speech_to_text_model(model): - return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location) + return VertexChirpRealtimeConfig( + resolve_access_token=resolve_access_token, project=project, location=location, drop_params=drop_params + ) return VertexAIRealtimeConfig( access_token=access_token, project=project, diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py index 5efe125ee60..ffbf814788b 100644 --- a/litellm/llms/xai/realtime/transformation.py +++ b/litellm/llms/xai/realtime/transformation.py @@ -16,6 +16,7 @@ construction time (see ``handler.py``) so all normalization is isolated here and ``RealTimeStreaming`` stays provider-agnostic. """ +from collections.abc import Mapping from typing import Final @@ -77,6 +78,18 @@ class XAIRealtimeNormalizer: event = self._normalize_response_usage_event(event, event_type) return event + def gate_turn_detection(self, _turn_detection: object, /) -> Mapping[str, object]: + """xAI auto-answers under ``server_vad`` even with ``create_response`` off, so the gate is always untyped.""" + return {"create_response": False} + + @property + def deletes_conversation_items(self) -> bool: + """xAI answers ``conversation.item.delete`` with ``Item not found`` until it commits the audio item. + + The guardrail blocks on an interim transcript, before that commit. + """ + return False + def patch_outgoing_session(self, session: dict) -> dict: """Patch a client ``session.update`` payload before forwarding to xAI. 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..1851e818ec7 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, @@ -70917,6 +71184,38 @@ "notes": "Self-hosted decision model; infrastructure costs are paid separately" } }, + "nace/drex-v1.0": { + "input_cost_per_token": 4e-08, + "litellm_provider": "nace", + "max_input_tokens": 32768, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v1.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v2.0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-latest": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -75218,5 +75517,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/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 9c8890efd2b..4090666e16f 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1345,7 +1345,9 @@ class MCPRequestHandler: if auth is None: return None refreshed = auth.model_copy() - if auth.mcp_admitted_user_subject and auth.user_id: + if (auth.mcp_admitted_user_subject or auth.is_session_token) and auth.user_id: + # A session credential (``lite login`` or a redeemed ID-JAG) has no key row, its authority is + # sealed in the token: refresh the live user policy and keep the token's team, routes and models. current = await MCPRequestHandler.reload_admitted_user(auth.user_id, requires_fresh_policy=True) refreshed.object_permission = current.object_permission refreshed.object_permission_id = current.object_permission_id diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 6e5b5486696..449adaa73e9 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + JWT_BEARER_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -71,6 +72,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( revoke_refresh_token, supported_grant_types, ) +from litellm.proxy._experimental.mcp_server.id_jag_grant import ( + basic_client_id, + jwt_bearer_available, + redeem_id_jag_assertion, +) from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( exchange_idp_subject_token, token_exchange_available, @@ -2396,7 +2402,7 @@ async def token_endpoint( grant_type: str = Form(...), code: str = Form(None), redirect_uri: str = Form(None), - client_id: str = Form(...), + client_id: str | None = Form(None), client_secret: str | None = Form(None), code_verifier: str = Form(None), refresh_token: str | None = Form(None), @@ -2405,6 +2411,7 @@ async def token_endpoint( subject_token: str | None = Form(None), subject_token_type: str | None = Form(None), requested_token_type: str | None = Form(None), + assertion: str | None = Form(None), mcp_server_name: str | None = None, ): """ @@ -2416,7 +2423,14 @@ async def token_endpoint( 3. Return the token 4. Return a virtual key in this response """ - if mcp_server_name is None and is_gateway_dcr_client_id(client_id): + presented_client_id: Final = client_id or basic_client_id(request.headers.get("authorization")) + if presented_client_id is None: + raise HTTPException( + status_code=400, detail={"error": "invalid_request", "error_description": "client_id is required"} + ) + if mcp_server_name is None and ( + is_gateway_dcr_client_id(presented_client_id) or grant_type == JWT_BEARER_GRANT_TYPE + ): from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load master_key, user_api_key_cache, @@ -2427,7 +2441,7 @@ async def token_endpoint( grant_type=grant_type, code=code, redirect_uri=redirect_uri, - client_id=client_id, + client_id=presented_client_id, code_verifier=code_verifier, refresh_token=refresh_token, master_key=master_key, @@ -2439,9 +2453,12 @@ async def token_endpoint( subject_token_type=subject_token_type, requested_token_type=requested_token_type, exchange_subject_token=exchange_idp_subject_token, + assertion=assertion, + client_secret=client_secret, + redeem_assertion=redeem_id_jag_assertion, ) - lookup_name: Final = mcp_server_name or client_id + lookup_name: Final = mcp_server_name or presented_client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if mcp_server is None and mcp_server_name is None: @@ -2454,7 +2471,7 @@ async def token_endpoint( grant_type=grant_type, code=code, redirect_uri=redirect_uri, - client_id=client_id, + client_id=presented_client_id, client_secret=client_secret, code_verifier=code_verifier, refresh_token=refresh_token, @@ -2568,7 +2585,8 @@ async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other language) reads to sign a user in through the browser and obtain a proxy credential.""" return JSONResponse( - native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS + native_client_auth_contract(request, token_exchange_available(), jwt_bearer_available()), + headers=TOKEN_NO_CACHE_HEADERS, ) @@ -3135,7 +3153,7 @@ def _aggregate_header_scopes(request: Request) -> tuple[str, ...]: def _build_aggregate_authorization_server_response( - request: Request, token_exchange_available: bool + request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False ) -> dict[str, object]: """RFC 8414 metadata for the gateway as the aggregate authorization server. @@ -3155,9 +3173,13 @@ def _build_aggregate_authorization_server_response( "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": supported_grant_types(token_exchange_available), + "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available), "code_challenge_methods_supported": ["S256"], - "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], + "token_endpoint_auth_methods_supported": ( + ["none", "client_secret_post", "client_secret_basic"] + if jwt_bearer_available + else ["none", "client_secret_post"] + ), } @@ -3193,7 +3215,7 @@ async def oauth_authorization_server_aggregate(request: Request): per-server row win here instead would serve an issuer of {base} against a resource that advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - return _build_aggregate_authorization_server_response(request, token_exchange_available()) + return _build_aggregate_authorization_server_response(request, token_exchange_available(), jwt_bearer_available()) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} @@ -3273,6 +3295,10 @@ def _build_oauth_authorization_server_response( if explicitly_named else request_base_url ) + # With no upstream resolved the document is the gateway's own: its issuer is the bare origin + # and its /token serves every grant the gateway does, including the ID-JAG jwt-bearer grant. + gateway_owned: Final = mcp_server_name is None and issuer_path is None + jwt_bearer: Final = gateway_owned and jwt_bearer_available() return { "issuer": issuer, @@ -3280,9 +3306,15 @@ def _build_oauth_authorization_server_response( "token_endpoint": token_endpoint, "response_types_supported": ["code"], "scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []), - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": ( + list(supported_grant_types(token_exchange_available(), jwt_bearer)) + if gateway_owned + else ["authorization_code", "refresh_token"] + ), "code_challenge_methods_supported": ["S256"], - "token_endpoint_auth_methods_supported": ["client_secret_post"], + "token_endpoint_auth_methods_supported": ( + ["none", "client_secret_post", "client_secret_basic"] if jwt_bearer else ["client_secret_post"] + ), # Claude expects a registration endpoint, even if we just fake it "registration_endpoint": ( f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register" diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index dd8c82a99fe..90289fd4c20 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -198,20 +198,26 @@ class MintProxyCredential(Protocol): and mints the same credential ``lite login`` mints.""" def __call__( - self, user_id: str, team_id: str | None, / + self, user_id: str, team_id: str | None, /, *, allowed_routes: Sequence[str] | None = None ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ... TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" -def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]: +def supported_grant_types(token_exchange_available: bool, jwt_bearer_available: bool = False) -> tuple[str, ...]: """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed only where the JWT auth that proves a subject token is on, backed by a database, and - licensed, so a client never selects a grant the gateway would then refuse.""" - if token_exchange_available: - return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE) - return ("authorization_code", "refresh_token") + licensed, so a client never selects a grant the gateway would then refuse. The RFC 7523 + jwt-bearer grant (an ID-JAG from the identity provider) is listed only where, on top of + that, the operator has allowlisted at least one requesting app.""" + return ( + "authorization_code", + "refresh_token", + *((TOKEN_EXCHANGE_GRANT_TYPE,) if token_exchange_available else ()), + *((JWT_BEARER_GRANT_TYPE,) if jwt_bearer_available else ()), + ) """RFC 8693: a native client that already holds a token from the customer's identity @@ -236,10 +242,29 @@ class SubjectIdentity(LiteLLMBaseModel): class SubjectTokenRefusal(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"] + error: Literal[ + "unsupported_grant_type", + "invalid_request", + "invalid_grant", + "invalid_client", + "invalid_scope", + "temporarily_unavailable", + ] description: str = Field(min_length=1) +class AssertionIdentity(SubjectIdentity): + """The litellm user an RFC 7523 assertion (an ID-JAG) stands for, plus what the grant + tail needs from the assertion itself: the single-use key for its ``jti``, how long that + marker must outlive it, the ``scope`` the credential is issued for, and the route groups that + scope narrows the credential to (``None`` leaves it the full ``lite login`` credential).""" + + replay_key: str = Field(min_length=1) + replay_ttl_seconds: int = Field(gt=0) + scope: str | None = None + allowed_routes: tuple[str, ...] | None = None + + class ExchangeSubjectToken(Protocol): """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the IdP token the way the proxy's own JWT auth does and names the litellm user and team it @@ -272,7 +297,9 @@ class LookupConsentTeams(Protocol): def __call__(self, user_id: str, /) -> Awaitable[tuple[ConsentTeam, ...] | ReloadUserFailure]: ... -async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCredentialMintFailure: +async def _refuse_proxy_credential( + user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None +) -> ProxyCredentialMintFailure: return "unresolvable" @@ -282,6 +309,25 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject ) +class RedeemAssertion(Protocol): + """Injected RFC 7523 assertion verifier ``(assertion, client_id, client_secret, request)``: + proves an ID-JAG the identity provider minted for the requesting app ``client_id``, + authenticates that app, and names the litellm user the assertion stands for, or says why + this gateway will not take it.""" + + def __call__( + self, assertion: str, client_id: str, client_secret: str | None, request: Request, / + ) -> Awaitable[AssertionIdentity | SubjectTokenRefusal]: ... + + +async def _refuse_assertion( + assertion: str, client_id: str, client_secret: str | None, request: Request +) -> SubjectTokenRefusal: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="this gateway is not configured to redeem identity assertions" + ) + + async def _unavailable_vendor_credential( user_id: str, server_id: str, requested_scopes: tuple[str, ...] = () ) -> VendorCredentialState: @@ -797,7 +843,9 @@ class NativeClientAuthContract(TypedDict): revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]] -def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract: +def native_client_auth_contract( + request: Request, token_exchange_available: bool, jwt_bearer_available: bool = False +) -> NativeClientAuthContract: """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a native client (in any language) needs to run the sign-in without reading LiteLLM source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter @@ -812,7 +860,7 @@ def native_client_auth_contract(request: Request, token_exchange_available: bool "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": supported_grant_types(token_exchange_available), + "grant_types_supported": supported_grant_types(token_exchange_available, jwt_bearer_available), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), @@ -1403,10 +1451,11 @@ class _ProxyCredentialTokenResponse(TypedDict): access_token: ReadOnly[str] token_type: ReadOnly[Literal["Bearer"]] expires_in: ReadOnly[int] - refresh_token: ReadOnly[str] + refresh_token: NotRequired[ReadOnly[str]] user_id: ReadOnly[str] team_id: ReadOnly[str | None] issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]] + scope: NotRequired[ReadOnly[str]] def proxy_credential_response( @@ -1415,6 +1464,7 @@ def proxy_credential_response( signing: SessionSigning, issued_token_type: _IssuedTokenType | None = None, family: str | None = None, + scope: str | None = None, ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and @@ -1426,7 +1476,7 @@ def proxy_credential_response( refresh: Final = mint_session_refresh_token(bound_principal, signing.keys, signing.now, family=family) if not isinstance(refresh, MintedSessionToken): return _oauth_error(500, "server_error", "failed to mint the session credential") - credential: Final[_ProxyCredentialTokenResponse] = { + renewable: Final[_ProxyCredentialTokenResponse] = { "access_token": minted.key, "token_type": "Bearer", "expires_in": minted.expires_in, @@ -1434,12 +1484,27 @@ def proxy_credential_response( "user_id": minted.user_id, "team_id": minted.team_id, } + credential: Final[_ProxyCredentialTokenResponse] = renewable if scope is None else {**renewable, "scope": scope} if issued_token_type is None: return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type} return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS) +def assertion_credential_response(minted: MintedProxyCredential, scope: str | None) -> Response: + """The proxy-API token response for an assertion grant: the same access token, with the + scope the assertion carried, and no refresh token.""" + unscoped: Final[_ProxyCredentialTokenResponse] = { + "access_token": minted.key, + "token_type": "Bearer", + "expires_in": minted.expires_in, + "user_id": minted.user_id, + "team_id": minted.team_id, + } + credential: Final[_ProxyCredentialTokenResponse] = unscoped if scope is None else {**unscoped, "scope": scope} + return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) + + def _reload_failure_response(failure: ReloadUserFailure) -> Response: """Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new ``ReloadUserFailure`` member is a type error here rather than silently 400ing.""" @@ -1460,7 +1525,9 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response: match refusal.error: case "temporarily_unavailable": return _oauth_error(503, refusal.error, refusal.description) - case "unsupported_grant_type" | "invalid_request": + case "invalid_client": + return _oauth_error(401, refusal.error, refusal.description) + case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope": return _oauth_error(400, refusal.error, refusal.description) case _: assert_never(refusal.error) @@ -1515,6 +1582,9 @@ async def aggregate_token( subject_token_type: str | None = None, requested_token_type: str | None = None, exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token, + assertion: str | None = None, + client_secret: str | None = None, + redeem_assertion: RedeemAssertion = _refuse_assertion, ) -> Response: """The aggregate token verb: authorization_code and refresh_token grants for the identity-only session pair, or for the proxy-API credential when the grant was issued @@ -1561,10 +1631,19 @@ async def aggregate_token( exchange_subject_token=exchange_subject_token, issue=issue, ) + if grant_type == JWT_BEARER_GRANT_TYPE: + return await _jwt_bearer_grant( + request=request, + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + redeem_assertion=redeem_assertion, + issue=issue, + ) return _oauth_error( 400, "unsupported_grant_type", - f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}", + f"grant_type must be authorization_code, refresh_token, {TOKEN_EXCHANGE_GRANT_TYPE}, or {JWT_BEARER_GRANT_TYPE}", ) @@ -1690,6 +1769,48 @@ class _GrantIssuer: return _mint_failure_response(minted) return proxy_credential_response(minted, principal, self._signing, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE) + async def redeem_assertion( + self, assertion: str, client_id: str, client_secret: str | None, redeem_assertion: RedeemAssertion + ) -> Response: + """The RFC 7523 tail: prove the assertion and the requesting app, claim the assertion's + ``jti`` before anything is minted (RFC 7523 section 3 makes a replayed ``jti`` a + rejection, and the identity provider mints a fresh assertion cheaply, so a mint that + fails after the claim costs the client one more round trip, never a replay window), + then mint the proxy credential the user is revalidated for. No refresh token comes + with it (RFC 7521 section 4.1): the app renews by asking its identity provider for a + fresh assertion, which keeps the user's standing there in the loop.""" + target_refusal: Final = self._assertion_target_refusal() + if target_refusal is not None: + return target_refusal + identity: Final = await redeem_assertion(assertion, client_id, client_secret, self._request) + if isinstance(identity, SubjectTokenRefusal): + return _subject_token_refusal_response(identity) + refused: Final = _claim_refusal( + await self._guard.claim(identity.replay_key, identity.replay_ttl_seconds), + _oauth_error(400, "invalid_grant", "the assertion has already been redeemed"), + ) + if refused is not None: + return refused + minted: Final = await self._mint_proxy_credential( + identity.user_id, identity.team_id, allowed_routes=identity.allowed_routes + ) + if not isinstance(minted, MintedProxyCredential): + return _mint_failure_response(minted) + return assertion_credential_response(minted, identity.scope) + + def _assertion_target_refusal(self) -> Response | None: + """An ID-JAG names the proxy origin or its MCP resource (``{base}/mcp``) as ``aud``, and an RFC + 8707 client sends that same value as ``resource``, so both are this grant's targets.""" + if self._resource is None or is_proxy_api_resource(self._request, self._resource): + return None + if self._resource.rstrip("/") == f"{get_request_base_url(self._request)}/mcp": + return None + return _oauth_error( + 400, + "invalid_target", + "resource does not match the proxy API or MCP resource this gateway issues credentials for", + ) + def _proxy_api_target_refusal(self) -> Response | None: if self._resource is None or is_proxy_api_resource(self._request, self._resource): return None @@ -1807,6 +1928,19 @@ async def refresh_proxy_credential( ) +async def _jwt_bearer_grant( + request: Request, + assertion: str | None, + client_id: str, + client_secret: str | None, + redeem_assertion: RedeemAssertion, + issue: _GrantIssuer, +) -> Response: + if not assertion: + return _oauth_error(400, "invalid_request", "assertion is required") + return await issue.redeem_assertion(assertion, client_id, client_secret, redeem_assertion) + + async def _token_exchange_grant( subject_token: str | None, subject_token_type: str | None, diff --git a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py new file mode 100644 index 00000000000..155d5280f5a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py @@ -0,0 +1,540 @@ +"""The RFC 7523 jwt-bearer grant on ``POST /token``: a requesting app that obtained an +ID-JAG (Identity Assertion Authorization Grant, the cross-app access profile) from the +customer's identity provider redeems it for the same proxy-API credential ``lite login`` +stores. The assertion is proven with the JWT auth issuer and JWKS plumbing the proxy +already trusts, its claims are checked the way the ID-JAG draft prescribes, the requesting +app must be one the operator allowlisted, and the ``jti`` is single-use until the +assertion expires.""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import hmac +import math +from collections.abc import Awaitable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timezone +from functools import partial +from types import MappingProxyType +from typing import Final, Protocol +from urllib.parse import unquote_plus + +import jwt +from fastapi import Request + +from litellm._logging import verbose_proxy_logger +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + AssertionIdentity, + ConsentTeam, + SubjectIdentity, + SubjectTokenRefusal, +) +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + gateway_outage_refusal, + log_line, + read_token_exchange_prerequisites, +) +from litellm.proxy._experimental.mcp_server.oauth_utils import proxy_base_url_from_env +from litellm.proxy._types import IdJagRequestingApp +from litellm.proxy.auth.handle_jwt import ID_JAG_TYP, AssertionRejection, JWTHandler + +ID_JAG_REQUIRED_CLAIMS: Final = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id") +REJECTED_ASSERTION: Final = "assertion was rejected" +TEAM_HEADER: Final = "x-litellm-team-id" +_INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable")) +_REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:" +SCOPE_MCP: Final = "litellm:mcp" +SCOPE_LLM: Final = "litellm:llm" +SCOPE_ROUTE_GROUPS: Final = MappingProxyType( + {SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)} +) + + +@dataclass(frozen=True, slots=True) +class GrantedScope: + """What an assertion's ``scope`` claim buys: the scope echoed in the token response, exactly the + tokens enforced, and the route groups the credential is narrowed to. No claim means the full + credential, so both are ``None``.""" + + scope: str | None + allowed_routes: tuple[str, ...] | None + + +def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal: + """The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and + no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity + provider named that resource, so the credential is ceilinged to the MCP routes whatever the + claim says: ``litellm:llm`` cannot be granted and is refused.""" + if claim is not None and not isinstance(claim, str): + return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) + requested: Final = frozenset(claim.split()) if claim else frozenset() + if not requested <= SCOPE_ROUTE_GROUPS.keys(): + return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) + if mcp_resource: + if SCOPE_LLM in requested: + return SubjectTokenRefusal(error="invalid_scope", description=_MCP_RESOURCE_DESCRIPTION) + return GrantedScope(scope=SCOPE_MCP, allowed_routes=SCOPE_ROUTE_GROUPS[SCOPE_MCP]) + if not requested: + return GrantedScope(scope=None, allowed_routes=None) + granted: Final = tuple(scope for scope in SCOPE_ROUTE_GROUPS if scope in requested) + return GrantedScope(scope=" ".join(granted), allowed_routes=tuple(_route_groups(granted))) + + +def names_mcp_resource(aud: object, base_url: str | None) -> bool: + """Whether the verified ``aud`` names only the MCP resource, ``{base}/mcp``, and not the + authorization server itself (the bare origin).""" + audiences: Final = frozenset((aud,) if isinstance(aud, str) else aud if isinstance(aud, list) else ()) + return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences + + +def _route_groups(scopes: Sequence[str]) -> Iterator[str]: + for scope in scopes: + yield from SCOPE_ROUTE_GROUPS[scope] + + +_SCOPE_DESCRIPTION: Final = ( + f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated" +) +_MCP_RESOURCE_DESCRIPTION: Final = ( + f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes" +) + + +@dataclass(frozen=True, slots=True) +class IdJagPrerequisites: + """The deployment-level gates for redeeming an ID-JAG: the JWT auth whose issuer plumbing + proves the assertion must be on, backed by a database, and licensed, not mapping IdP tokens to + virtual keys (whose blocks and budgets a minted credential would sidestep), with a shared Redis + so an assertion is redeemed once across every worker and replica (the in-memory marker is per + process, so a second worker would mint again), ``PROXY_BASE_URL`` + must fix the gateway URL the assertion's ``aud`` has to name (a request's ``Host`` is the + caller's to choose), and the operator must have allowlisted the requesting apps. Discovery advertises the grant only when every gate + holds, and a redemption is refused naming the first one that does not.""" + + jwt_auth_enabled: bool + has_database: bool + licensed: bool + maps_jwts_to_virtual_keys: bool + shared_replay_store: bool + base_url: str | None + requesting_apps: tuple[IdJagRequestingApp, ...] + + @property + def available(self) -> bool: + return ( + self.jwt_auth_enabled + and self.has_database + and self.licensed + and not self.maps_jwts_to_virtual_keys + and self.shared_replay_store + and self.base_url is not None + and bool(self.requesting_apps) + ) + + def refusal(self) -> SubjectTokenRefusal | None: + if not self.jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot redeem identity assertions", + ) + if not self.has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot redeem identity assertions", + ) + if not self.licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is an enterprise only feature; no license is set", + ) + if self.maps_jwts_to_virtual_keys: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway maps IdP tokens to virtual keys, which the jwt-bearer grant does not serve", + ) + if not self.shared_replay_store: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="configure Redis (cache_params or coordination_redis) so this gateway can redeem an identity assertion exactly once across its workers", + ) + if self.base_url is None: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="set PROXY_BASE_URL to this gateway's public URL, the audience an identity assertion must name", + ) + if not self.requesting_apps: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no requesting apps allowlisted for identity assertions", + ) + return None + + +def read_id_jag_prerequisites() -> IdJagPrerequisites: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + jwt_handler, + redis_usage_cache, + ) + + shared: Final = read_token_exchange_prerequisites() + return IdJagPrerequisites( + jwt_auth_enabled=shared.jwt_auth_enabled, + has_database=shared.has_database, + licensed=shared.licensed, + maps_jwts_to_virtual_keys=shared.maps_jwts_to_virtual_keys, + shared_replay_store=redis_usage_cache is not None, + base_url=proxy_base_url_from_env(), + requesting_apps=_requesting_apps(jwt_handler), + ) + + +def _requesting_apps(jwt_handler: JWTHandler) -> tuple[IdJagRequestingApp, ...]: + if not hasattr(jwt_handler, "litellm_jwtauth"): + return () + return tuple(jwt_handler.litellm_jwtauth.id_jag_requesting_apps or ()) + + +def jwt_bearer_available() -> bool: + return read_id_jag_prerequisites().available + + +def gateway_audiences(base_url: str | None) -> tuple[str, ...]: + """The issuer values this gateway answers to, which an ID-JAG's ``aud`` must name: its MCP + resource and its authorization server, both under the fixed ``PROXY_BASE_URL``.""" + return () if base_url is None else (f"{base_url}/mcp", base_url) + + +class DecodeAssertion(Protocol): + """Injected signature and claim verification ``(assertion, audience, required_claims)``: + ``JWTHandler.decode_trusted_assertion`` in production, which returns the verified claims, + a rejection, or raises when the gateway could not reach the issuer's keys.""" + + def __call__( + self, assertion: str, audience: Sequence[str], required_claims: Sequence[str], / + ) -> Awaitable[Mapping[str, object] | AssertionRejection]: ... + + +class ClaimReader(Protocol): + """The JWT auth user mapping as ``JWTHandler`` exposes it: the litellm user id and email the + verified claims name through ``user_id_jwt_field`` and ``user_email_jwt_field``.""" + + def get_user_id(self, token: dict[str, object], default_value: str | None) -> str | None: ... + + def get_user_email(self, token: dict[str, object], default_value: str | None) -> str | None: ... + + def get_all_jwt_team_ids(self, token: dict[str, object]) -> list[str]: ... + + def is_allowed_domain(self, user_email: str) -> bool: ... + + +class ResolveSubject(Protocol): + """Injected subject lookup ``(claims, team_header)``: the litellm user the verified claims + name through the JWT auth user mapping (``user_id_jwt_field`` or ``user_email_jwt_field``, + else ``sub``) and the team the credential is minted for (see ``select_team``), ``None`` when + the gateway knows no such user, or a refusal when the team cannot be settled. Raises on a + database outage.""" + + def __call__( + self, claims: Mapping[str, object], team_header: str | None, / + ) -> Awaitable[SubjectIdentity | SubjectTokenRefusal | None]: ... + + +class LookupUser(Protocol): + """Injected ``(user_id, user_email)`` lookup: the stored id of the existing litellm user the + claims name, ``None`` when the gateway knows no such user. Raises on a database outage.""" + + def __call__(self, user_id: str, user_email: str | None, /) -> Awaitable[str | None]: ... + + +class LookupTeams(Protocol): + """Injected lookup of the teams a user is on, or the failure text when they cannot be read.""" + + def __call__(self, user_id: str, /) -> Awaitable[Sequence[ConsentTeam] | str]: ... + + +class GetUser(Protocol): + def __call__( + self, *, user_id: str, user_email: str | None, sso_user_id: str + ) -> Awaitable[LiteLLM_UserTable | None]: ... + + +def proxy_user_lookup(get_user: GetUser) -> LookupUser: + """``LookupUser`` over the proxy's user table. The assertion's subject is matched as the + local ``user_id`` and as the ``sso_user_id`` an SSO login linked, the way the JWT and SSO + login paths match it, so a linked account needs no email claim.""" + + async def lookup_user(user_id: str, user_email: str | None) -> str | None: + user: Final = await get_user(user_id=user_id, user_email=user_email, sso_user_id=user_id) + return None if user is None else user.user_id + + return lookup_user + + +def subject_resolver(reader: ClaimReader, lookup_user: LookupUser, lookup_teams: LookupTeams) -> ResolveSubject: + """``ResolveSubject`` over the JWT auth claim mapping and the proxy's user and team lookups.""" + + async def resolve_subject( + claims: Mapping[str, object], team_header: str | None + ) -> SubjectIdentity | SubjectTokenRefusal | None: + token: Final = dict(claims) + subject: Final = claims.get("sub") + claimed_user_id: Final = reader.get_user_id(token=token, default_value=None) or ( + subject if isinstance(subject, str) else None + ) + if claimed_user_id is None: + return None + user_email: Final = reader.get_user_email(token=token, default_value=None) + if not reader.is_allowed_domain(user_email or ""): + return SubjectTokenRefusal( + error="invalid_grant", + description="the assertion's subject email is missing or outside the allowed domain", + ) + user_id: Final = await lookup_user(claimed_user_id, user_email) + if user_id is None: + return None + teams: Final = await lookup_teams(user_id) + if teams in _INACTIVE_SUBJECT: + return SubjectTokenRefusal( + error="invalid_grant", description="the assertion's subject is not an active user" + ) + if isinstance(teams, str): + return SubjectTokenRefusal( + error="temporarily_unavailable", description="the gateway could not read this user's teams; retry" + ) + team: Final = select_team(team_header, tuple(reader.get_all_jwt_team_ids(token=token)), teams) + if isinstance(team, SubjectTokenRefusal): + return team + return SubjectIdentity(user_id=user_id, team_id=team) + + return resolve_subject + + +def select_team( + team_header: str | None, claim_teams: Sequence[str], teams: Sequence[ConsentTeam] +) -> str | None | SubjectTokenRefusal: + """The team an ID-JAG redemption mints for, settled before the assertion is consumed since + the minter refuses a teamless credential for a user who has a team. The ``x-litellm-team-id`` + header (a team id or alias) wins and must name one of the user's teams, then the assertion's + mapped team claim, then the user's only team. A user on several teams must name one.""" + if team_header is not None: + named: Final = next((team for team in teams if team_header in (team.team_id, team.team_alias)), None) + if named is None: + return SubjectTokenRefusal( + error="invalid_request", description=f"{TEAM_HEADER} does not name one of this user's teams" + ) + return named.team_id + asserted: Final = tuple(team.team_id for team in teams if team.team_id in claim_teams) + if len(asserted) == 1: + return asserted[0] + if len(asserted) > 1: + return SubjectTokenRefusal( + error="invalid_request", + description=f"the assertion names several of this user's teams; send the {TEAM_HEADER} header naming the team for this credential", + ) + if len(teams) == 1: + return teams[0].team_id + if teams: + return SubjectTokenRefusal( + error="invalid_request", + description=f"this user belongs to several teams; send the {TEAM_HEADER} header naming the team for this credential", + ) + return None + + +@dataclass(frozen=True, slots=True) +class AssertionRedemption: + """One ``jwt-bearer`` request as the grant sees it: the assertion, how the requesting app + identified and authenticated itself, the team the credential is asked for, the issuer + values this gateway answers to (what the assertion's ``aud`` must name), the clock, and the + leeway the verifier grants ``exp``, so the replay marker outlives every moment the assertion + would still verify.""" + + assertion: str + client_id: str + client_secret: str | None + authorization_header: str | None + team_id: str | None + audiences: tuple[str, ...] + now: float + leeway_seconds: float + + +async def redeem_id_jag_assertion( + assertion: str, client_id: str, client_secret: str | None, request: Request +) -> AssertionIdentity | SubjectTokenRefusal: + from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( # noqa: PLC0415 # circular import at module load + lookup_consent_teams_from_database, + ) + from litellm.proxy.auth.auth_checks import get_user_object # noqa: PLC0415 # circular import at module load + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + jwt_handler, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + resolve_subject: Final = subject_resolver( + _claim_reader(jwt_handler), + proxy_user_lookup( + partial( + get_user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + ), + lookup_consent_teams_from_database, + ) + prerequisites: Final = read_id_jag_prerequisites() + return await identity_from_id_jag( + AssertionRedemption( + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + authorization_header=request.headers.get("authorization"), + team_id=request.headers.get(TEAM_HEADER), + audiences=gateway_audiences(prerequisites.base_url), + now=datetime.now(timezone.utc).timestamp(), + leeway_seconds=jwt_handler.leeway, + ), + prerequisites=prerequisites, + decode=jwt_handler.decode_trusted_assertion, + resolve_subject=resolve_subject, + ) + + +def _claim_reader(handler: ClaimReader) -> ClaimReader: + return handler + + +async def identity_from_id_jag( + redemption: AssertionRedemption, + prerequisites: IdJagPrerequisites, + decode: DecodeAssertion, + resolve_subject: ResolveSubject, +) -> AssertionIdentity | SubjectTokenRefusal: + """Prove an ID-JAG for the grant tail. The requesting app is checked first, the way a token + endpoint authenticates its client before it looks at the grant, so an app the operator never + allowlisted learns nothing about the assertion. Everything wrong with the assertion itself is + ``invalid_grant`` (RFC 7523 section 3.1), a check the gateway could not complete is + ``temporarily_unavailable``, and the reason stays in the proxy log because this endpoint is + public and the verifier's wording can quote the assertion.""" + unmet: Final = prerequisites.refusal() + if unmet is not None: + return unmet + app: Final = next((app for app in prerequisites.requesting_apps if app.client_id == redemption.client_id), None) + if app is None: + return SubjectTokenRefusal(error="invalid_client", description="client_id is not an allowlisted requesting app") + if not _client_authenticated(app, redemption.client_secret, redemption.authorization_header): + return SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + try: + header: Final = jwt.get_unverified_header(redemption.assertion) + except jwt.PyJWTError: + return SubjectTokenRefusal(error="invalid_grant", description="assertion is not a JWT") + if header.get("typ") != ID_JAG_TYP: + return SubjectTokenRefusal(error="invalid_grant", description=f"assertion typ must be {ID_JAG_TYP}") + try: + verified: Final = await decode(redemption.assertion, redemption.audiences, ID_JAG_REQUIRED_CLAIMS) + except Exception as denied: # noqa: BLE001 # the verifier raises its 503 for a JWKS outage, and anything else is a bad assertion + return gateway_outage_refusal(denied, denied, what="an assertion", parameter="assertion") or _rejected( + str(denied) + ) + if isinstance(verified, AssertionRejection): + return _rejected(verified.reason) + marker: Final = _replay_marker(verified, redemption) + if isinstance(marker, SubjectTokenRefusal): + return marker + replay_key, replay_ttl_seconds = marker + scope: Final = granted_scope( + verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url) + ) + if isinstance(scope, SubjectTokenRefusal): + return scope + try: + subject: Final = await resolve_subject(verified, redemption.team_id) + except Exception as failed: # noqa: BLE001 # the user lookup raises for a missing user and for a database outage alike + return gateway_outage_refusal( + failed, failed, what="an assertion's subject", parameter="assertion" + ) or _rejected(f"subject lookup failed: {failed}") + if subject is None: + return _rejected("the assertion names no user the gateway knows") + if isinstance(subject, SubjectTokenRefusal): + return subject + return AssertionIdentity( + user_id=subject.user_id, + team_id=subject.team_id, + replay_key=replay_key, + replay_ttl_seconds=replay_ttl_seconds, + scope=scope.scope, + allowed_routes=scope.allowed_routes, + ) + + +def _replay_marker( + verified: Mapping[str, object], redemption: AssertionRedemption +) -> tuple[str, int] | SubjectTokenRefusal: + """The single-use key and its lifetime for a verified assertion, once the assertion is bound to + the presenting app, carries no proof-of-possession the gateway cannot check, and is still alive.""" + if verified.get("client_id") != redemption.client_id: + return _rejected("the assertion's client_id is not the requesting app that presented it") + if "cnf" in verified: + return _rejected("the assertion is proof-bound and this gateway does not check proof of possession") + issuer: Final = verified.get("iss") + jti: Final = verified.get("jti") + exp: Final = verified.get("exp") + if not isinstance(issuer, str) or not isinstance(jti, str) or not jti or not isinstance(exp, (int, float)): + return _rejected("the assertion's iss, jti, or exp claim has the wrong shape") + replay_ttl_seconds: Final = math.ceil(exp + redemption.leeway_seconds - redemption.now) + if replay_ttl_seconds <= 0: + return _rejected("the assertion has expired") + return _REPLAY_KEY_PREFIX + hashlib.sha256(f"{issuer}\n{jti}".encode()).hexdigest(), replay_ttl_seconds + + +def _rejected(reason: str) -> SubjectTokenRefusal: + verbose_proxy_logger.warning("jwt-bearer grant refused an assertion: %s", log_line(reason)) + return SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + + +def _client_authenticated(app: IdJagRequestingApp, form_secret: str | None, authorization_header: str | None) -> bool: + """A public app (no secret configured) is identified by ``client_id`` alone. A confidential + app must present its secret, by HTTP Basic (RFC 6749 section 2.3.1, the credentials + form-urlencoded inside the base64) or the ``client_secret`` form field, compared in constant + time.""" + if app.client_secret is None: + return True + basic: Final = _basic_credentials(authorization_header) + if basic is not None: + return basic[0] == app.client_id and _same_secret(basic[1], app.client_secret) + return form_secret is not None and _same_secret(form_secret, app.client_secret) + + +def _same_secret(presented: str, expected: str) -> bool: + return hmac.compare_digest(presented.encode(), expected.encode()) + + +def basic_client_id(authorization_header: str | None) -> str | None: + """The ``client_id`` an HTTP Basic ``Authorization`` header carries, for a confidential app + that leaves the form field out as RFC 6749 section 2.3.1 allows.""" + credentials: Final = _basic_credentials(authorization_header) + return None if credentials is None else credentials[0] + + +def _basic_credentials(authorization_header: str | None) -> tuple[str, str] | None: + if authorization_header is None: + return None + scheme, _, encoded = authorization_header.strip().partition(" ") + if scheme.lower() != "basic" or not encoded: + return None + try: + decoded: Final = base64.b64decode(encoded.strip(), validate=True).decode() + except (binascii.Error, UnicodeDecodeError): + return None + client_id, separator, secret = decoded.partition(":") + if not separator: + return None + return unquote_plus(client_id), unquote_plus(secret) diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 15632eb4783..5044a5bd3dc 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -20,13 +20,15 @@ from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" -SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = ( - "the gateway could not verify subject_token because its identity provider or database is unavailable; retry" +_CHECK_UNAVAILABLE: Final = ( + "the gateway could not verify {parameter} because its identity provider or database is unavailable; retry" ) -SUBJECT_TOKEN_CHECK_FAULTED: Final = ( - "the gateway could not verify subject_token because its database reported a fault that is not a transient " +_CHECK_FAULTED: Final = ( + "the gateway could not verify {parameter} because its database reported a fault that is not a transient " "outage; retrying will not help until the gateway deployment is repaired" ) +SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = _CHECK_UNAVAILABLE.format(parameter="subject_token") +SUBJECT_TOKEN_CHECK_FAULTED: Final = _CHECK_FAULTED.format(parameter="subject_token") GatewayOutage = Literal["retryable", "faulted"] @@ -182,20 +184,37 @@ async def identity_from_subject_token( def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal: + outage: Final = gateway_outage_refusal(denied, reason, what="a subject_token", parameter="subject_token") + if outage is not None: + return outage + verbose_proxy_logger.warning("token exchange refused a subject_token: %s", log_line(reason)) + return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + + +def gateway_outage_refusal(denied: Exception, reason: object, what: str, parameter: str) -> SubjectTokenRefusal | None: + """The ``temporarily_unavailable`` refusal when ``denied`` is the gateway failing to check + ``what`` (its database or the IdP's JWKS down), naming the grant ``parameter`` it could not verify, + else ``None`` because the token itself was bad.""" outage: Final = _gateway_could_not_verify(denied) if outage is None: - verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) - return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) - verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason) - return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage)) + return None + verbose_proxy_logger.error("token endpoint could not verify %s, %s: %s", what, outage, log_line(reason)) + return SubjectTokenRefusal( + error="temporarily_unavailable", description=_check_unavailable_description(outage, parameter) + ) -def _check_unavailable_description(outage: GatewayOutage) -> str: +def log_line(reason: object) -> str: + """``reason`` as one log line, so text a caller controls cannot forge log entries.""" + return str(reason).replace("\r", " ").replace("\n", " ") + + +def _check_unavailable_description(outage: GatewayOutage, parameter: str) -> str: match outage: case "retryable": - return SUBJECT_TOKEN_CHECK_UNAVAILABLE + return _CHECK_UNAVAILABLE.format(parameter=parameter) case "faulted": - return SUBJECT_TOKEN_CHECK_FAULTED + return _CHECK_FAULTED.format(parameter=parameter) case _: assert_never(outage) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fe34e18aecd..b724a61cd94 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -455,6 +455,7 @@ class MCPServerConfig(TypedDict, total=False): upstream_token_header: ReadOnly[str] id_jag_resource_token_endpoint: str id_jag_resource: str + id_jag_subject_source: ReadOnly[str] id_jag_resource_client_id: ReadOnly[str] id_jag_resource_client_secret: ReadOnly[str] client_private_key: str @@ -2789,6 +2790,7 @@ class MCPServerManager: # ID-JAG fields id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None), id_jag_resource=server_config.get("id_jag_resource", None), + id_jag_subject_source=server_config.get("id_jag_subject_source", None), id_jag_resource_client_id=server_config.get("id_jag_resource_client_id", None), id_jag_resource_client_secret=server_config.get("id_jag_resource_client_secret", None), client_private_key=server_config.get("client_private_key", None), @@ -3380,6 +3382,7 @@ class MCPServerManager: credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None ), id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None), + id_jag_subject_source=(credentials_dict.get("id_jag_subject_source") if credentials_dict else None), id_jag_resource_client_id=(credentials_dict.get("id_jag_resource_client_id") if credentials_dict else None), id_jag_resource_client_secret=self._decrypt_credential_field( credentials_dict.get("id_jag_resource_client_secret") if credentials_dict else None, diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b0448a62c9e..133bad780bf 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -19,6 +19,9 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( build_token_endpoint_client_auth, normalize_token_endpoint_auth_method, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + GOOGLE_REFRESH_TOKEN_PARAMS, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path @@ -59,7 +62,7 @@ _DEFAULT_NATIVE_REDIRECT_URIS: Final[list[str]] = [ ] _GOOGLE_AUTHORIZATION_HOSTS: Final = frozenset({"accounts.google.com"}) -_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"}) +_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = GOOGLE_REFRESH_TOKEN_PARAMS _warned_invalid_proxy_base_url: str | None = None @@ -128,6 +131,11 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st return urlunparse(parsed._replace(query=urlencode(merged))) +def proxy_base_url_from_env() -> str | None: + """The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed.""" + return _resolve_proxy_base_url_env() + + def _resolve_proxy_base_url_env() -> str | None: global _warned_invalid_proxy_base_url configured: Final = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 2de54e8ce1b..90f01368127 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -1160,6 +1160,7 @@ async def _get_tools_from_mcp_servers( protocol_version: str | None = None, *, record_listing: bool = False, + request_trace_id: str | None = None, enforce_rate_limits: bool = True, ) -> AggregateToolListing: """ @@ -1206,6 +1207,7 @@ async def _get_tools_from_mcp_servers( "spend_logs_metadata": spend_logs_metadata, "headers": logging_safe_mcp_headers(raw_headers), **({"tags": request_tags} if request_tags else {}), + **({"existing_trace_id": request_trace_id} if request_trace_id else {}), }, # Provide a small input payload for standard logging "input": [ diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py index 20fac28a55b..3260603c3ca 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -39,6 +39,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( HeaderCarrier, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -71,6 +72,7 @@ __all__ = [ "HeaderCarrier", "IdJagConfig", "IdJagResourceClient", + "IdJagSubjectSource", "NoOpAuth", "NoneConfig", "Ok", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 935d026b22c..f217072bd39 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -34,6 +34,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( CredError, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -236,7 +237,8 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None: if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None: return None resource_client: Final = _id_jag_resource_client(server, client_id, client_auth) - if resource_client is None: + subject_source: Final = _id_jag_subject_source(server) + if resource_client is None or subject_source is None: return None return ServerSpec( server_id=server.server_id, @@ -249,6 +251,7 @@ def _id_jag_spec(server: MCPServer, resource: str) -> ServerSpec | None: client_auth=client_auth, resource_client=resource_client, subject_token_type=_id_jag_subject_token_type(server), + subject_source=subject_source, audience=server.audience, resource=server.id_jag_resource, scopes=tuple(server.scopes or ()), @@ -282,6 +285,18 @@ def _id_jag_client_auth(server: MCPServer) -> ClientAuth | None: return None +def _id_jag_subject_source(server: MCPServer) -> IdJagSubjectSource | None: + """``None`` for a value that names no source, so the half-configured server defers like a missing + endpoint does instead of silently presenting the wrong token.""" + configured: Final = server.id_jag_subject_source + if not configured: + return IdJagSubjectSource.login_assertion + try: + return IdJagSubjectSource(configured) + except ValueError: + return None + + def _id_jag_subject_token_type(server: MCPServer) -> str: """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token; an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim.""" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py new file mode 100644 index 00000000000..fcd8ae22ce3 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py @@ -0,0 +1,264 @@ +"""Exchange the identity assertion captured at SSO login for a refresh token, at login. + +An ``oauth2_id_jag`` server configured with ``id_jag_subject_source: refresh_token`` does not present the +login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued for that assertion through +one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at +an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at +login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and +the refresh token outlives the assertion by as long as the IdP decides. An OIDC login signed in through the very +IdP client the ID-JAG server uses already holds such a refresh token: it is bound to that client, so it is +presented as is and no further exchange is posted, which Okta refuses for an id_token subject anyway. + +The servers consulted are the two configuration authorities, the pod-local config declaration and the DB +rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a +server added on another pod. The exchange authenticates as the ID-JAG server's own IdP client, since the +refresh token is bound to the client that later spends it. One row holds one refresh token per user, so every server in this mode has +to share one IdP client; a second client is logged and skipped rather than silently handed a token it +cannot redeem. A refused or unreachable IdP never fails the login, the row keeps the plain assertion and +the ID-JAG call reports the missing refresh token as the precondition it is. +""" + +from __future__ import annotations + +import asyncio +import os +from collections.abc import Iterable, Mapping, Sequence +from datetime import datetime, timezone +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + SSOClientConfig, + sso_client_config, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SSOIdentityAssertion, + with_login_refresh_token, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + CredError, + IdJagConfig, + IdJagSubjectSource, + ServerSpec, +) +from litellm.proxy._types import LiteLLM_MCPServerTable +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +REFRESH_TOKEN_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token" +_LOGIN_EXCHANGE_DEADLINE_SECONDS: Final = 15.0 + +_REFRESH_TOKEN_SCOPE: Final = "openid offline_access" +_RUNTIME_APPROVAL_STATES: Final = (None, "active", "approved") # the states get_runtime_mcp_server_rows serves + + +class LoginExchangeCatalog(Protocol): + """The slice of the MCP server manager the login-time exchange reads.""" + + @property + def config_mcp_servers(self) -> Mapping[str, MCPServer]: ... + + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: ... + + +class TokenExchangeEndpoint(Protocol): + async def fetch( + self, endpoint: str, client_id: str, grant_params: Mapping[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: ... + + +def _is_runtime_id_jag_row(row: LiteLLM_MCPServerTable) -> bool: + """An ID-JAG row the runtime registry would serve. A pending or rejected submission never gets a + login assertion posted to the token endpoint it names.""" + return row.auth_type == MCPAuth.oauth2_id_jag and row.approval_status in _RUNTIME_APPROVAL_STATES + + +async def login_exchange_servers( + catalog: LoginExchangeCatalog, db_rows: Sequence[LiteLLM_MCPServerTable] +) -> tuple[MCPServer, ...]: + """Every server the exchange has to consider: the config-declared ones plus the approved ID-JAG rows + in the DB, built the same way the catalog builds them so their credentials are decrypted. A DB row + hides a config server with the same ``server_id``, exactly as the live catalog serves them.""" + id_jag_rows: Final = tuple(row for row in db_rows if _is_runtime_id_jag_row(row)) + db_servers: Final = tuple([await _built_or_skipped(catalog, row) for row in id_jag_rows]) + db_by_id: Final = {server.server_id: server for server in db_servers if server is not None} + return tuple({**catalog.config_mcp_servers, **db_by_id}.values()) + + +async def _built_or_skipped(catalog: LoginExchangeCatalog, row: LiteLLM_MCPServerTable) -> MCPServer | None: + """One row that cannot be built (undecryptable credentials, a half-configured submission) is skipped + with a warning. It must not cost every other refresh-token server its login exchange.""" + try: + return await catalog.build_mcp_server_from_table(row, register_oauth_discovery=False) + except Exception as exc: # noqa: BLE001 # a per-row build failure is logged and isolated, not propagated + verbose_proxy_logger.warning( + "ID-JAG: MCP server row server_id=%s could not be built for the login-time refresh token exchange " + "and is skipped: %s", + row.server_id, + exc, + ) + return None + + +def refresh_token_exchange_targets(specs: Iterable[ServerSpec | None]) -> tuple[IdJagConfig, ...]: + """The ID-JAG configs that present a login refresh token, one per IdP client, in registration order.""" + configs: Final = tuple( + spec.config + for spec in specs + if spec is not None + and isinstance(spec.config, IdJagConfig) + and spec.config.subject_source is IdJagSubjectSource.refresh_token + ) + clients: Final = tuple((config.org_token_endpoint, config.client_id) for config in configs) + return tuple(config for index, config in enumerate(configs) if clients.index(clients[index]) == index) + + +async def exchange_login_assertion_for_refresh_token( + assertion: SSOIdentityAssertion, + targets: tuple[IdJagConfig, ...], + token_endpoint: TokenExchangeEndpoint, + now: datetime, + deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS, + sso_client: SSOClientConfig | None = None, +) -> SSOIdentityAssertion: + """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the + refresh token one target's IdP client can redeem. An OIDC login's own refresh token is that token + when the SSO client that received it is the target's client at the same token endpoint; otherwise it + is bound to a client the ID-JAG client is not, and the assertion is exchanged for one at the first + target's client.""" + if not targets or assertion.exchanged_refresh_token is not None: + return assertion + reusable: Final = _target_redeeming_the_login_refresh_token(assertion, targets, sso_client) + if reusable is not None and assertion.refresh_token is not None: + verbose_proxy_logger.info( + "ID-JAG: the refresh token issued at OIDC login belongs to client_id=%s at %s, the IdP client of a server " + "with id_jag_subject_source: refresh_token, so it is kept as that subject and no further exchange is posted", + reusable.client_id, + reusable.org_token_endpoint, + ) + return with_login_refresh_token( + assertion, + assertion.refresh_token.get_secret_value(), + None, + now, + endpoint=reusable.org_token_endpoint, + client_id=reusable.client_id, + ) + config: Final = targets[0] + if len(targets) > 1: + verbose_proxy_logger.warning( + "ID-JAG: %d IdP clients are configured with id_jag_subject_source refresh_token; the login refresh " + "token is minted for client_id=%s at %s and servers using another client cannot redeem it.", + len(targets), + config.client_id, + config.org_token_endpoint, + ) + params: Final = { + "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, + "requested_token_type": REFRESH_TOKEN_TOKEN_TYPE, + "subject_token": assertion.id_token.get_secret_value(), + "subject_token_type": assertion.subject_token_type or config.subject_token_type, + "scope": _REFRESH_TOKEN_SCOPE, + } + try: + outcome: Final = await asyncio.wait_for( + token_endpoint.fetch(config.org_token_endpoint, config.client_id, params, config.client_auth), + timeout=deadline_seconds, + ) + except (TimeoutError, asyncio.TimeoutError): + verbose_proxy_logger.warning( + "ID-JAG: the IdP at %s did not answer the login-time refresh token exchange within %.0fs, so the login " + "finishes with the captured assertion only. Servers with id_jag_subject_source: refresh_token will ask " + "this user to sign in again.", + config.org_token_endpoint, + deadline_seconds, + ) + return assertion + match outcome: + case Error(err): + verbose_proxy_logger.warning( + "ID-JAG: the IdP at %s refused to exchange the login assertion for a refresh token (%s); the " + "login goes on with the assertion alone and servers presenting a refresh token will ask the " + "user to sign in again.", + config.org_token_endpoint, + err.summary, + ) + return assertion + case Ok(token): + verbose_proxy_logger.info( + "ID-JAG: exchanged the login assertion for a refresh token at %s (expires_in=%s)", + config.org_token_endpoint, + token.expires_in, + ) + return with_login_refresh_token( + assertion, + token.access_token, + token.expires_in, + now, + endpoint=config.org_token_endpoint, + client_id=config.client_id, + ) + + +def _target_redeeming_the_login_refresh_token( + assertion: SSOIdentityAssertion, targets: tuple[IdJagConfig, ...], sso_client: SSOClientConfig | None +) -> IdJagConfig | None: + """The target whose IdP client is the SSO client the login's refresh token was issued to. A refresh + token is redeemable only by the client that received it at the endpoint that issued it, so both have + to match exactly before the login token is presented anywhere.""" + if assertion.refresh_token is None or sso_client is None: + return None + return next( + ( + config + for config in targets + if config.org_token_endpoint == sso_client.token_endpoint and config.client_id == sso_client.client_id + ), + None, + ) + + +async def exchange_login_assertion_if_configured( + assertion: SSOIdentityAssertion | None, +) -> SSOIdentityAssertion | None: + """The SSO-callback hook: consults the configured and DB-declared MCP servers and performs the + exchange when one of them presents a login refresh token. Nothing here can fail the login.""" + if assertion is None: + return None + from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers # noqa: PLC0415 # import cycle + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids import cycle + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # import cycle + to_server_spec, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( # noqa: PLC0415 # cycle + TokenEndpointClient, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global + + try: + db_rows: Final = () if prisma_client is None else await get_all_mcp_servers(prisma_client) + servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows) + targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers) + return await exchange_login_assertion_for_refresh_token( + assertion, + targets, + TokenEndpointClient(), + datetime.now(timezone.utc), + sso_client=sso_client_config(os.environ), + ) + except Exception as exc: # noqa: BLE001 # the login itself must not fail on an egress-side exchange + verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc) + return assertion diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 66b44e91fd7..ea0cdee5c1c 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -20,6 +20,7 @@ v2: no imports from v1. from __future__ import annotations import hashlib +from dataclasses import dataclass from datetime import datetime, timezone from functools import partial from typing import Final @@ -53,7 +54,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, SSOAssertionStore, + SSOIdentityAssertion, assertion_expired, + exchanged_refresh_token_issued_for, + login_refresh_token_expired, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -75,6 +79,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ClientSecretAuth, CredError, IdJagConfig, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -88,6 +93,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( _TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" _JWT_BEARER_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:jwt-bearer" _ID_JAG_REQUESTED_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:id-jag" +_REFRESH_TOKEN_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:refresh_token" + + +@dataclass(frozen=True, slots=True) +class _IdJagSubjectToken: + """What ID-JAG leg 1 asserts and the RFC 8693 type it is declared as.""" + + value: str + token_type: str class _NullOAuthTokenStore: @@ -200,13 +214,15 @@ class UpstreamCredentialProvider: async def _id_jag( self, subject: Subject, server: ServerSpec, config: IdJagConfig ) -> Result[httpx2.Auth, CredError]: - match await self._id_jag_subject_token(subject): + match await self._id_jag_subject_token(subject, config): case Error(err): return Error(err) case Ok(subject_token): return await self._id_jag_exchange(subject, subject_token, server, config) - async def _id_jag_subject_token(self, subject: Subject) -> Result[str, CredError]: + async def _id_jag_subject_token( + self, subject: Subject, config: IdJagConfig + ) -> Result[_IdJagSubjectToken, CredError]: """The identity token ID-JAG leg 1 asserts, from the request or from the SSO login it was captured at. A caller that presents its own IdP identity token wins: that is the strongest available @@ -223,7 +239,7 @@ class UpstreamCredentialProvider: a raised error so a DB blip cannot 500 the egress or the upstream-401 retry. """ if subject.inbound_token is not None: - return Ok(subject.inbound_token.get_secret_value()) + return Ok(_IdJagSubjectToken(subject.inbound_token.get_secret_value(), config.subject_token_type)) if not subject.subject_id: return Error( CredError.of_precondition_required( @@ -232,7 +248,11 @@ class UpstreamCredentialProvider: ) ) try: - assertion: Final = await self._sso_assertion_store.fetch(subject.subject_id) + assertion: Final = await ( + self._sso_assertion_store.fetch_stored(subject.subject_id) + if config.subject_source is IdJagSubjectSource.refresh_token + else self._sso_assertion_store.fetch(subject.subject_id) + ) except AssertionStoreUnavailable as exc: # The driver's message can name hosts, schemas or connection details, and this summary # is returned to the caller verbatim as a 503 body. Operators get it from the log. @@ -253,6 +273,8 @@ class UpstreamCredentialProvider: "Sign in through LiteLLM SSO so the gateway captures one." ) ) + if config.subject_source is IdJagSubjectSource.refresh_token: + return _login_refresh_token_subject(assertion, config, datetime.now(timezone.utc)) if assertion_expired(assertion, datetime.now(timezone.utc)): return Error( CredError.of_precondition_required( @@ -260,20 +282,24 @@ class UpstreamCredentialProvider: "LiteLLM SSO again to capture a current one." ) ) - return Ok(assertion.id_token.get_secret_value()) + return Ok( + _IdJagSubjectToken( + assertion.id_token.get_secret_value(), assertion.subject_token_type or config.subject_token_type + ) + ) async def _id_jag_exchange( - self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig + self, subject: Subject, token: _IdJagSubjectToken, server: ServerSpec, config: IdJagConfig ) -> Result[httpx2.Auth, CredError]: slot: Final = _id_jag_slot_key(subject, server) - fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config) + fingerprint: Final = _id_jag_fingerprint(token.value, server.server_id, config) async def _exchange() -> Result[ExchangedToken, CredError]: leg1_params: Final = { "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE, - "subject_token": token, - "subject_token_type": config.subject_token_type, + "subject_token": token.value, + "subject_token_type": token.token_type, **({"audience": config.audience} if config.audience else {}), **({"resource": config.resource} if config.resource else {}), **({"scope": " ".join(config.scopes)} if config.scopes else {}), @@ -399,6 +425,42 @@ class UpstreamCredentialProvider: return None +def _login_refresh_token_subject( + assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime +) -> Result[_IdJagSubjectToken, CredError]: + """The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh + token the ID-JAG client holds from the user's sign-in, exchanged from the login assertion or, when + the SSO client is this server's IdP client, the one the OIDC login itself returned. The login-time + hook records which client and endpoint it is bound to, and a token bound elsewhere is never sent + here. The assertion's own expiry does not apply, a SAML assertion is spent minutes after login + while the refresh token it bought lives as long as the IdP said.""" + if assertion.exchanged_refresh_token is None: + return Error( + CredError.of_precondition_required( + "This server mints its ID-JAG from a refresh token exchanged at SSO login and none is stored " + "for this user. Sign in through LiteLLM SSO again; if the login log shows the IdP refused the " + "exchange, the IdP client has to allow a refresh token for this requesting app." + ) + ) + if not exchanged_refresh_token_issued_for(assertion, config.org_token_endpoint, config.client_id): + return Error( + CredError.of_precondition_required( + "The refresh token exchanged at SSO login was minted for a different IdP client than this server " + "uses, so it cannot be presented here. Point every id_jag_subject_source: refresh_token server at " + "the same IdP client, then sign in through LiteLLM SSO again." + ) + ) + if login_refresh_token_expired(assertion, now): + return Error( + CredError.of_precondition_required( + "The refresh token exchanged at SSO login for this user has expired. Sign in through LiteLLM SSO again." + ) + ) + return Ok( + _IdJagSubjectToken(assertion.exchanged_refresh_token.get_secret_value(), _REFRESH_TOKEN_SUBJECT_TOKEN_TYPE) + ) + + def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str: """Which cache slot this caller's bearer for this upstream lives in. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py index bca5848febe..ae5e9d0417e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -12,15 +12,18 @@ not the size of the user table) and single-flighted per user through the same `` the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent redemptions of one token as replay and can revoke the whole grant chain. -The refresh is redeemed against the generic-OIDC client the login itself used -(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy -reconciles from the stored SSO row into the process environment at startup), authenticated the way -that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the -credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept -only one of the two. An assertion can only exist if that client minted it, so no other client could -redeem its refresh token, and no other method is known to be accepted. A deployment whose -``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss -logs the scope by name rather than failing silently. +The refresh is redeemed against the SSO client the login itself used, resolved in the same order +the login callback dispatches (Google, then Microsoft, then generic OIDC, from the client ids the +proxy reconciles from the stored SSO row into the process environment at startup), authenticated the +way that login authenticated: the Google and Microsoft logins and the non-PKCE generic path send +HTTP Basic, while the PKCE generic path sends the credentials in the body when +``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept only one of the two. An +assertion can only exist if that client minted it, so no other client could redeem its refresh +token, and no other method is known to be accepted. A login whose authorization request never asked +for a refresh token captures none at all: the Google and Microsoft logins ask (``access_type=offline``, +``offline_access``) only while an ``oauth2_id_jag`` server exists, and generic OIDC asks only when +``GENERIC_SCOPE`` carries ``offline_access``, which is why that miss names the cause rather than +failing silently. Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged @@ -77,7 +80,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s SSOIdentityAssertion, assertion_expired, assertion_from_sso_login, + carry_exchanged_refresh_token, fetch_sso_identity_assertion, + is_saml_assertion, persist_sso_identity_assertion, ) from litellm.types.llms.custom_http import httpxSpecialProvider @@ -92,6 +97,10 @@ _SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion" # Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of # the exchange is replaced first. Matches the sibling per-user token store's skew. _DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0 +# Where fastapi-sso's Google login redeemed the authorization code, read from Google's discovery +# document; Microsoft's mirrors ``CustomMicrosoftSSO``, override first, tenant template second. +_GOOGLE_TOKEN_ENDPOINT: Final = "https://oauth2.googleapis.com/token" +_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token" class AssertionRead(Protocol): @@ -122,7 +131,7 @@ class FormPost(Protocol): @dataclass(frozen=True, slots=True) class SSOClientConfig: - """The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how.""" + """The SSO client credentials a refresh_token grant has to authenticate as, and how.""" token_endpoint: str client_id: str @@ -131,32 +140,58 @@ class SSOClientConfig: def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None: - """The configured generic-OIDC client, or ``None`` when the deployment has none. + """The configured SSO client the login used, or ``None`` when the deployment has none. - Read from the process environment because that is where the login path reads it - (``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row - at startup, so this resolves to the same client that minted the assertion. ``None`` is an - ordinary state, not an error: a deployment signing in through a provider that captures no + Read from the process environment because that is where the login callback reads it and where + the proxy materializes the stored ``sso_config`` row at startup, and resolved in the callback's + own order (``GOOGLE_CLIENT_ID``, then ``MICROSOFT_CLIENT_ID``, then ``GENERIC_CLIENT_ID``, each + selecting its branch when set at all), so this is the same client that minted the assertion: a + deployment carrying both a Google and a generic client signed the user in through Google. ``None`` + is an ordinary state, not an error: a deployment signing in through a provider that captures no assertion has nothing here to renew, and a client with no secret is not a confidential client that could redeem one. - ``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because - the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while - the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so - a constant here would authenticate the renewal differently from the sign-in that produced the + ``auth_method`` follows how each login authenticated at the token endpoint. The Google and + Microsoft logins (fastapi-sso) and the non-PKCE generic path send HTTP Basic, while the PKCE + generic path puts the credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set. A + constant here would authenticate the renewal differently from the sign-in that produced the refresh token and 401 against an IdP application registered for only one of the two. """ - token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT") - client_id: Final = env.get("GENERIC_CLIENT_ID") - client_secret: Final = env.get("GENERIC_CLIENT_SECRET") + if env.get("GOOGLE_CLIENT_ID") is not None: + return _confidential_client( + _GOOGLE_TOKEN_ENDPOINT, env.get("GOOGLE_CLIENT_ID"), env.get("GOOGLE_CLIENT_SECRET"), "client_secret_basic" + ) + if env.get("MICROSOFT_CLIENT_ID") is not None: + tenant: Final = env.get("MICROSOFT_TENANT") + return _confidential_client( + env.get("MICROSOFT_TOKEN_ENDPOINT") + or (_MICROSOFT_TOKEN_ENDPOINT_TEMPLATE.format(tenant=tenant) if tenant else None), + env.get("MICROSOFT_CLIENT_ID"), + env.get("MICROSOFT_CLIENT_SECRET"), + "client_secret_basic", + ) + includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" + return _confidential_client( + env.get("GENERIC_TOKEN_ENDPOINT"), + env.get("GENERIC_CLIENT_ID"), + env.get("GENERIC_CLIENT_SECRET"), + "client_secret_post" if includes_client_id else "client_secret_basic", + ) + + +def _confidential_client( + token_endpoint: str | None, + client_id: str | None, + client_secret: str | None, + auth_method: MCPTokenEndpointAuthMethod, +) -> SSOClientConfig | None: if not token_endpoint or not client_id or not client_secret: return None - includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" return SSOClientConfig( token_endpoint=token_endpoint, client_id=client_id, client_secret=SecretStr(client_secret), - auth_method="client_secret_post" if includes_client_id else "client_secret_basic", + auth_method=auth_method, ) @@ -260,11 +295,19 @@ class SSOAssertionRefresher: async def refresh( self, user_id: str, assertion: SSOIdentityAssertion ) -> Result[SSOIdentityAssertion, RefreshFailure]: + if is_saml_assertion(assertion): + verbose_proxy_logger.warning( + "ID-JAG: the stored SAML assertion for user_id=%s has reached its NotOnOrAfter. A SAML assertion " + "cannot be renewed, so the user has to sign in through SAML SSO again.", + user_id, + ) + return Error(RefreshFailure.of_rejected("a SAML assertion cannot be renewed; sign in again")) if assertion.refresh_token is None: verbose_proxy_logger.warning( "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was " - "captured with it, so it cannot be renewed without another interactive sign-in. Add " - "'offline_access' to GENERIC_SCOPE so the SSO login captures one.", + "captured with it, so it cannot be renewed without another interactive sign-in. Google and Microsoft " + "logins ask for one once an oauth2_id_jag MCP server is registered, so a user who signed in before " + "that has to sign in again. Generic OIDC asks only with 'offline_access' in GENERIC_SCOPE.", user_id, ) return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in")) @@ -272,11 +315,12 @@ class SSOAssertionRefresher: if config is None: verbose_proxy_logger.warning( "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed " - "because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, " - "GENERIC_CLIENT_SECRET).", + "because the active SSO provider has no confidential client configured (GOOGLE_CLIENT_ID and " + "GOOGLE_CLIENT_SECRET, MICROSOFT_CLIENT_ID with MICROSOFT_CLIENT_SECRET and MICROSOFT_TENANT, or " + "GENERIC_TOKEN_ENDPOINT with GENERIC_CLIENT_ID and GENERIC_CLIENT_SECRET).", user_id, ) - return Error(RefreshFailure.of_rejected("the generic SSO client is not configured")) + return Error(RefreshFailure.of_rejected("the SSO client is not configured")) carried_refresh_token: Final = assertion.refresh_token.get_secret_value() # Whichever method the SSO login used for this client, since that is the one the IdP @@ -323,10 +367,11 @@ class SSOAssertionRefresher: user_id, ) return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token")) - failure: Final = await self._store_renewal(user_id, previous, renewed) + carried: Final = carry_exchanged_refresh_token(previous, renewed) + failure: Final = await self._store_renewal(user_id, previous, carried) if failure is not None: return Error(failure) - return Ok(renewed) + return Ok(carried) async def _store_renewal( self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion @@ -401,9 +446,12 @@ class RefreshingSSOAssertionStore: self._skew = timedelta(seconds=expiry_skew_seconds) self._clock = clock + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self._inner.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: assertion: Final = await self._inner.fetch(user_id) - if not self._expiring(assertion): + if assertion is None or is_saml_assertion(assertion) or not self._expiring(assertion): return assertion await self._coordinator().run( user_id, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 0bc17d94165..19b3bc294d5 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -1,9 +1,9 @@ """Store for the enterprise IdP identity assertion captured at SSO login (EMA). -The ``oauth2_id_jag`` egress arm needs the user's IdP ``id_token`` as its RFC 8693 -``subject_token``. A front-door client holds an identity-only ``llm_session_`` bearer, not an -IdP assertion, so the assertion captured at the one SSO login is the only usable subject -source for it. This module owns both sides of that state: the SSO callback persists here +The ``oauth2_id_jag`` egress arm needs the user's IdP identity assertion as its RFC 8693 +``subject_token``: the OIDC ``id_token``, or the signed SAML assertion when the login was SAML. +A front-door client holds an identity-only ``llm_session_`` bearer, not an IdP assertion, so +the assertion captured at the one SSO login is the only usable subject source for it. This module owns both sides of that state: the SSO callback persists here (write-through to the DB so a login on one pod is visible to every pod) and the resolver seam reads back by ``user_id``. Retention is gated on an ``oauth2_id_jag`` server actually being registered, so a gateway with no EMA upstream never stores bearer material. @@ -17,9 +17,11 @@ TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against st from __future__ import annotations +import base64 import json from collections.abc import Mapping, Sequence -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol import jwt @@ -35,7 +37,15 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +"""What a Google or Microsoft authorization request has to carry for the IdP to issue a refresh token +beside the id_token. Asked for only while an ``oauth2_id_jag`` server exists to renew an assertion for, so +a deployment without one sees no extra consent.""" +MICROSOFT_REFRESH_TOKEN_SCOPE: Final = "offline_access" +GOOGLE_REFRESH_TOKEN_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"}) + + _ASSERTION_DECRYPT_LOG_KEY: Final = "sso_identity_assertion" +SAML2_SUBJECT_TOKEN_TYPE: Final = "urn:ietf:params:oauth:token-type:saml2" _STR_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _MAYBE_STR_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) @@ -69,8 +79,10 @@ def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable: class SSOIdentityAssertion(LiteLLMBaseModel): - """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token, - ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login.""" + """The IdP material an EMA exchange needs: ``id_token`` is the RFC 8693 subject token (the OIDC + id_token, or the base64 signed ```` when ``subject_token_type`` names SAML), + ``expires_at`` bounds its usefulness, and the refresh token renews it without re-login. + ``subject_token_type`` is ``None`` when the exchange should send the operator-configured type.""" model_config = ConfigDict(frozen=True) @@ -78,6 +90,11 @@ class SSOIdentityAssertion(LiteLLMBaseModel): refresh_token: SecretStr | None = None issuer: str | None = None expires_at: datetime | None = None + subject_token_type: str | None = None + exchanged_refresh_token: SecretStr | None = None + exchanged_refresh_token_expires_at: datetime | None = None + exchanged_refresh_token_endpoint: str | None = None + exchanged_refresh_token_client_id: str | None = None class SSOAssertionCache: @@ -130,6 +147,11 @@ class _StoredAssertionPayload(LiteLLMBaseModel): refresh_token: str | None = None issuer: str | None = None expires_at: datetime | None = None + subject_token_type: str | None = None + exchanged_refresh_token: str | None = None + exchanged_refresh_token_expires_at: datetime | None = None + exchanged_refresh_token_endpoint: str | None = None + exchanged_refresh_token_client_id: str | None = None def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIdentityAssertion | None: @@ -159,6 +181,94 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden ) +def assertion_from_saml_login( + assertion_xml: bytes, not_on_or_after: int | None, issuer: str | None +) -> SSOIdentityAssertion: + """The carrier for a SAML login: the signed ```` element the IdP issued, base64 + encoded the way RFC 8693 carries a ``saml2`` subject token, with its own validity window as + the expiry and no refresh token, since SAML has none and the user has to sign in again.""" + return SSOIdentityAssertion( + id_token=SecretStr(base64.urlsafe_b64encode(assertion_xml).rstrip(b"=").decode("ascii")), + issuer=issuer, + expires_at=datetime.fromtimestamp(not_on_or_after, tz=timezone.utc) if not_on_or_after is not None else None, + subject_token_type=SAML2_SUBJECT_TOKEN_TYPE, + ) + + +def is_saml_assertion(assertion: SSOIdentityAssertion) -> bool: + return assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + + +def with_login_refresh_token( + assertion: SSOIdentityAssertion, + refresh_token: str, + expires_in: int | None, + now: datetime, + *, + endpoint: str | None = None, + client_id: str | None = None, +) -> SSOIdentityAssertion: + """The same login, now also carrying the refresh token the ID-JAG client exchanged the assertion + for, with the lifetime the IdP gave that token. It is kept apart from ``refresh_token``, which an + OIDC login holds for the SSO client and which only that client can redeem. The assertion and its + own expiry stay as they were, so a server that presents the login assertion directly is unaffected.""" + return assertion.model_copy( + update={ + "exchanged_refresh_token": SecretStr(refresh_token), + "exchanged_refresh_token_expires_at": now + timedelta(seconds=expires_in) + if expires_in is not None + else None, + "exchanged_refresh_token_endpoint": endpoint, + "exchanged_refresh_token_client_id": client_id, + } + ) + + +def carry_exchanged_refresh_token( + previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion +) -> SSOIdentityAssertion: + """The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login. + Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to + them and stays valid for as long as the IdP said. When the login's own refresh token was kept as + that subject (the SSO client is the ID-JAG client), a rotation the renewal just received replaces + it too, since the IdP has retired the token the row held.""" + login_token_reused: Final = ( + previous.refresh_token is not None + and previous.exchanged_refresh_token is not None + and previous.exchanged_refresh_token.get_secret_value() == previous.refresh_token.get_secret_value() + ) + return renewed.model_copy( + update={ + "exchanged_refresh_token": renewed.refresh_token + if login_token_reused + else previous.exchanged_refresh_token, + "exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at, + "exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint, + "exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id, + } + ) + + +def exchanged_refresh_token_issued_for(assertion: SSOIdentityAssertion, endpoint: str, client_id: str) -> bool: + """Whether the exchanged refresh token was minted at this token endpoint for this client. A refresh + token is bound to the client that received it, so one minted for another server's IdP client must + never be posted to this server's endpoint.""" + return ( + assertion.exchanged_refresh_token_endpoint == endpoint + and assertion.exchanged_refresh_token_client_id == client_id + ) + + +def login_refresh_token_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: + """Whether the refresh token exchanged at login has passed the lifetime the IdP gave it. One + with no recorded lifetime is left for the IdP to judge, like an assertion with no ``exp``.""" + expires_at: Final = assertion.exchanged_refresh_token_expires_at + if expires_at is None: + return False + normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) + return normalized <= now + + def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: """Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is treated as usable and left for the IdP to reject, since the store records what the id_token @@ -212,6 +322,27 @@ async def persist_sso_identity_assertion( **({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}), **({"issuer": assertion.issuer} if assertion.issuer else {}), **({"expires_at": assertion.expires_at.isoformat()} if assertion.expires_at else {}), + **({"subject_token_type": assertion.subject_token_type} if assertion.subject_token_type else {}), + **( + {"exchanged_refresh_token": assertion.exchanged_refresh_token.get_secret_value()} + if assertion.exchanged_refresh_token + else {} + ), + **( + {"exchanged_refresh_token_expires_at": assertion.exchanged_refresh_token_expires_at.isoformat()} + if assertion.exchanged_refresh_token_expires_at + else {} + ), + **( + {"exchanged_refresh_token_endpoint": assertion.exchanged_refresh_token_endpoint} + if assertion.exchanged_refresh_token_endpoint + else {} + ), + **( + {"exchanged_refresh_token_client_id": assertion.exchanged_refresh_token_client_id} + if assertion.exchanged_refresh_token_client_id + else {} + ), } encoded: Final = _STR_ADAPTER.validate_python(encrypt_value_helper(json.dumps(payload))) await _assertion_table(prisma_client).upsert( @@ -250,6 +381,11 @@ async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None: refresh_token=SecretStr(payload.refresh_token) if payload.refresh_token else None, issuer=payload.issuer, expires_at=payload.expires_at, + subject_token_type=payload.subject_token_type, + exchanged_refresh_token=SecretStr(payload.exchanged_refresh_token) if payload.exchanged_refresh_token else None, + exchanged_refresh_token_expires_at=payload.exchanged_refresh_token_expires_at, + exchanged_refresh_token_endpoint=payload.exchanged_refresh_token_endpoint, + exchanged_refresh_token_client_id=payload.exchanged_refresh_token_client_id, ) @@ -290,6 +426,12 @@ class SSOAssertionStore(Protocol): async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: ... + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + """The row as stored, never renewed on the way out. For a server that presents the refresh + token exchanged at login, the SSO id_token's own lifetime is irrelevant and an SSO renewal + failing must not stand between that server and a still valid refresh token.""" + ... + class DbSSOAssertionStore: """The live store: the row the SSO callback wrote, read back by ``user_id``. @@ -308,6 +450,9 @@ class DbSSOAssertionStore: except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence raise AssertionStoreUnavailable(str(exc)) from exc + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self.fetch(user_id) + async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None: try: return await _read_assertion_from_db(user_id) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 71745a560f3..45c53af9349 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -349,6 +349,14 @@ class IdJagResourceClient(LiteLLMBaseModel): client_auth: ClientAuth +class IdJagSubjectSource(str, Enum): + """What ID-JAG leg 1 presents as the RFC 8693 subject_token: the identity assertion captured at SSO + login, or the refresh token that assertion was exchanged for at the IdP when the user logged in.""" + + login_assertion = "login_assertion" + refresh_token = "refresh_token" + + class IdJagConfig(HeaderCarrier): """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange"). @@ -368,6 +376,7 @@ class IdJagConfig(HeaderCarrier): client_auth: ClientAuth resource_client: IdJagResourceClient subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token" + subject_source: IdJagSubjectSource = IdJagSubjectSource.login_assertion audience: str | None = None resource: str | None = None scopes: tuple[str, ...] = () diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index 2f7fcaef645..fd637a2f187 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -25,7 +25,16 @@ from litellm.proxy.management_endpoints.ui_sso import ( async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: - user: Final = await load_active_user_by_id(user_id) + return await _consent_teams(await load_active_user_by_id(user_id)) + + +async def lookup_consent_teams_from_database(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: + """The same teams read from the database row, for a grant that is consumed once it is settled: the + membership the minter will re-check must be the membership the team was picked from.""" + return await _consent_teams(await load_active_user_by_id(user_id, source="database")) + + +async def _consent_teams(user: LiteLLM_UserTable | ReloadUserFailure) -> tuple[ConsentTeam, ...] | ReloadUserFailure: if isinstance(user, str): return user details: Final = await _team_details(user.teams) @@ -39,7 +48,7 @@ async def lookup_consent_teams(user_id: str) -> tuple[ConsentTeam, ...] | Reload async def mint_proxy_credential( - user_id: str, team_id: str | None + user_id: str, team_id: str | None, *, allowed_routes: Sequence[str] | None = None ) -> MintedProxyCredential | ProxyCredentialMintFailure: """Mint the ``lite login`` credential for a consented grant. Membership is checked live against the database row, so a team the user left between consent and redemption (or between refreshes) @@ -53,7 +62,9 @@ async def mint_proxy_credential( user row handed to the minter carries no team list, exactly like ``lite login``'s, so the minter's own first-team fallback stays inert. The credential carries the role the proxy already enforces for the user on every request, so a row with no role (JWT auth's - upsert writes none) mints as an internal user instead of being refused.""" + upsert writes none) mints as an internal user instead of being refused. ``allowed_routes`` + narrows the credential to those route groups, the way a virtual key's ``allowed_routes`` does, + for a grant whose scope asked for less than the full credential.""" user: Final = await load_active_user_by_id(user_id, source="database") if isinstance(user, str): return user @@ -75,6 +86,7 @@ async def mint_proxy_credential( team_alias=selected.team_alias, team_models=selected.team_models, team_model_aliases=selected.team_model_aliases, + allowed_routes=allowed_routes, ) return MintedProxyCredential( key=key, 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_features.py b/litellm/proxy/_lazy_features.py index a909a7a993a..a7adb5c1bf5 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -155,18 +155,13 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.mcp_management_endpoints", path_prefixes=("/v1/mcp/",), ), - LazyFeature( - # Also serves /.well-known/oauth-* (OAuth metadata discovery). - # No /mcp/oauth prefix here: the mounted /mcp sub-app would - # shadow it, and there are no actual routes there anyway. - name="mcp_byok_oauth", - module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", - path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), - ), LazyFeature( # Serves OAuth dance endpoints (/authorize, /token, /callback, # /register) plus several /.well-known/ discovery URLs at the proxy # root — needed for MCP-over-OAuth flows even before /mcp is hit. + # Registered before mcp_byok_oauth, which also claims the root + # /.well-known/oauth-authorization-server: the root document is the + # gateway's own authorization server metadata (token endpoint /token). name="mcp_discoverable", module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints", path_prefixes=( @@ -185,6 +180,14 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( # Catches the /{mcp_server_name}/authorize|token|register variants. path_suffixes=("/authorize", "/token", "/register"), ), + LazyFeature( + # Also serves /.well-known/oauth-* (OAuth metadata discovery). + # No /mcp/oauth prefix here: the mounted /mcp sub-app would + # shadow it, and there are no actual routes there anyway. + name="mcp_byok_oauth", + module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", + path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), + ), LazyFeature( name="mcp_rest", module_path="litellm.proxy._experimental.mcp_server.rest_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 20819a59a43..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": [ { @@ -34541,6 +34567,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -35771,7 +35808,7 @@ "/.well-known/oauth-authorization-server": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2", "parameters": [ { "in": "query", @@ -35819,7 +35856,7 @@ "/.well-known/oauth-authorization-server/mcp": { "get": { "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.", - "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get", + "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2", "responses": { "200": { "content": { @@ -35839,7 +35876,7 @@ "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}", - "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get", + "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -35880,7 +35917,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -35928,7 +35965,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": { "get": { "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.", - "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get", + "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2", "parameters": [ { "in": "path", @@ -35968,7 +36005,7 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", "responses": { "200": { "content": { @@ -36004,7 +36041,7 @@ "/.well-known/oauth-protected-resource/mcp": { "get": { "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.", - "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get", + "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2", "responses": { "200": { "content": { @@ -36024,7 +36061,7 @@ "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": { "get": { "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.", - "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get", + "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2", "parameters": [ { "in": "path", @@ -36065,7 +36102,7 @@ "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": { "get": { "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get", + "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2", "parameters": [ { "in": "path", @@ -36194,9 +36231,27 @@ }, "Body_token_endpoint__mcp_server_name__token_post": { "properties": { + "assertion": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Assertion" + }, "client_id": { - "title": "Client Id", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Id" }, "client_secret": { "anyOf": [ @@ -36293,17 +36348,34 @@ } }, "required": [ - "grant_type", - "client_id" + "grant_type" ], "title": "Body_token_endpoint__mcp_server_name__token_post", "type": "object" }, "Body_token_endpoint_token_post": { "properties": { + "assertion": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Assertion" + }, "client_id": { - "title": "Client Id", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Id" }, "client_secret": { "anyOf": [ @@ -36400,8 +36472,7 @@ } }, "required": [ - "grant_type", - "client_id" + "grant_type" ], "title": "Body_token_endpoint_token_post", "type": "object" @@ -37324,6 +37395,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -38151,7 +38233,7 @@ "/.well-known/oauth-authorization-server": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get_2", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server_get", "parameters": [ { "in": "query", @@ -38199,7 +38281,7 @@ "/.well-known/oauth-authorization-server/mcp": { "get": { "description": "OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414\npath-inserted form for a client that treats {base}/mcp as its authorization base URL.\n\nThe single-segment /mcp is reserved for the aggregate so the discovery chain stays\nconsistent: the aggregate protected-resource document advertises {base}/mcp as its\nauthorization server, so the document served here must have issuer {base}/mcp. A server\nliterally named ``mcp`` therefore does not take this route; it keeps its standard\ntwo-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the\nper-server row win here instead would serve an issuer of {base} against a resource that\nadvertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.", - "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get_2", + "operationId": "oauth_authorization_server_aggregate__well_known_oauth_authorization_server_mcp_get", "responses": { "200": { "content": { @@ -38219,7 +38301,7 @@ "/.well-known/oauth-authorization-server/mcp/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-authorization-server/mcp/{server_name}", - "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get_2", + "operationId": "oauth_authorization_server_mcp_standard__well_known_oauth_authorization_server_mcp__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38260,7 +38342,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}": { "get": { "description": "OAuth authorization server discovery endpoint.\n\nSupports both legacy pattern (/{server_name}) and root endpoint.", - "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get_2", + "operationId": "oauth_authorization_server_mcp__well_known_oauth_authorization_server__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38308,7 +38390,7 @@ "/.well-known/oauth-authorization-server/{mcp_server_name}/mcp": { "get": { "description": "OAuth authorization server discovery for legacy /{server_name}/mcp pattern.", - "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get_2", + "operationId": "oauth_authorization_server_legacy__well_known_oauth_authorization_server__mcp_server_name__mcp_get", "parameters": [ { "in": "path", @@ -38348,7 +38430,7 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", "responses": { "200": { "content": { @@ -38384,7 +38466,7 @@ "/.well-known/oauth-protected-resource/mcp": { "get": { "description": "OAuth protected resource discovery for the aggregate /mcp endpoint.\n\nThe single-segment ``/mcp`` path does not collide with any per-server PRM pattern\n(those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously\ndescribes the aggregate resource.", - "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get_2", + "operationId": "oauth_protected_resource_aggregate__well_known_oauth_protected_resource_mcp_get", "responses": { "200": { "content": { @@ -38404,7 +38486,7 @@ "/.well-known/oauth-protected-resource/mcp/{mcp_server_name}": { "get": { "description": "OAuth protected resource discovery endpoint using standard MCP URL pattern.\n\nStandard pattern: /mcp/{server_name}\nDiscovery path: /.well-known/oauth-protected-resource/mcp/{server_name}\n\nThis endpoint is compliant with MCP specification and works with standard\nMCP clients like mcp-inspector and VSCode Copilot.", - "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get_2", + "operationId": "oauth_protected_resource_mcp_standard__well_known_oauth_protected_resource_mcp__mcp_server_name__get", "parameters": [ { "in": "path", @@ -38445,7 +38527,7 @@ "/.well-known/oauth-protected-resource/{mcp_server_name}/mcp": { "get": { "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get_2", + "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource__mcp_server_name__mcp_get", "parameters": [ { "in": "path", @@ -40729,6 +40811,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { @@ -44809,6 +44902,17 @@ ], "title": "Id Jag Resource Token Endpoint" }, + "id_jag_subject_source": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Id Jag Subject Source" + }, "redirect_uris": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d62c6b0d51f..2c98cfeb33d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -551,6 +551,7 @@ class LiteLLMRoutes(enum.Enum): + list(agent_inference_routes) + model_info_routes ) + llm_api_routes_without_mcp = sorted(set(llm_api_routes) - set(mcp_inference_routes)) info_routes = [ "/key/info", "/key/health", @@ -4762,6 +4763,26 @@ class JWTIssuerConfig(LiteLLMBaseModel): DEFAULT_JWKS_STALE_TTL: Final = 3600 +class IdJagRequestingApp(LiteLLMBaseModel): + """A requesting application the operator lets redeem an ID-JAG (Identity Assertion + Authorization Grant) at the MCP gateway's ``POST /token``. ``client_id`` is the value + the identity provider puts in the assertion's ``client_id`` claim and the app sends as + the ``client_id`` form field. A ``client_secret`` makes the app a confidential client + that must authenticate with it (HTTP Basic or the ``client_secret`` form field); without + one the app is a public client.""" + + model_config = ConfigDict(frozen=True) + + client_id: str = Field( + min_length=1, description="The requesting app's client_id as the identity provider asserts it." + ) + client_secret: str | None = Field( + default=None, + min_length=1, + description="Shared secret the app must present at the gateway's token endpoint. Omit for a public client.", + ) + + class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): """ A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. @@ -4942,6 +4963,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=None, description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.", ) + id_jag_requesting_apps: tuple[IdJagRequestingApp, ...] | None = Field( + default=None, + description="Requesting apps allowed to redeem an ID-JAG at the MCP gateway's POST /token with the RFC 7523 jwt-bearer grant. Unset or empty turns that grant off.", + ) ######################################################### def __init__(self, **kwargs: Any) -> None: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index ee1ac79c820..1fa276d9ab2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -3,11 +3,13 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ 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 from litellm.anthropic_interface.exceptions import ( @@ -16,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, ) @@ -30,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 @@ -44,6 +48,8 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, 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() @@ -88,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. @@ -299,6 +330,26 @@ def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: return output_config.get("format") +_COUNTED_LITELLM_PARAMS: Final = ("additional_drop_params", "drop_params") + + +class _CountTokensBodyOptions(TypedDict): + output_format: ReadOnly[JsonValue] + context_management: ReadOnly[JsonValue] + request_litellm_params: ReadOnly[Mapping[str, JsonValue]] + + +def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: bool) -> _CountTokensBodyOptions: + sent_params: Final[Mapping[str, JsonValue]] = {key: data[key] for key in _COUNTED_LITELLM_PARAMS if key in data} + return { + "output_format": _requested_output_format(data), + "context_management": data.get("context_management"), + "request_litellm_params": MappingProxyType( + {"drop_params": True, **sent_params} if auto_drop_params else sent_params + ), + } + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -328,7 +379,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import count_request_tokens + 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: @@ -345,21 +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", - output_format=_requested_output_format(data), + **_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_checks.py b/litellm/proxy/auth/auth_checks.py index 8b32eb450e8..b395a826625 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3960,6 +3960,7 @@ class ExperimentalUIJWTToken: team_models: Sequence[str] | None = None, team_model_aliases: Mapping[str, str] | None = None, max_budget: float | None = None, + allowed_routes: Sequence[str] | None = None, ) -> str: """ Generate a JWT token for CLI authentication with configurable expiration. @@ -3973,6 +3974,8 @@ class ExperimentalUIJWTToken: team_alias: Team alias for the selected team, if available team_models: Model allowlist granted by the selected team team_model_aliases: Team model aliases for the selected team + allowed_routes: Route groups the session may call, enforced like a virtual key's + ``allowed_routes`` on every route; None leaves every route the role allows Returns: Encrypted JWT token string @@ -4017,6 +4020,7 @@ class ExperimentalUIJWTToken: max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), is_session_token=True, + allowed_routes=list(allowed_routes) if allowed_routes is not None else None, ) return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) 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/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 2421eb10c41..05b8cca19af 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -35,6 +35,7 @@ from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, RBAC_ROLES, + IdJagRequestingApp, JWKKeyValue, JWTAuthBuilderResult, JWTIssuerConfig, @@ -99,6 +100,10 @@ class NoMatchingJWTPublicKeyError(Exception): """Raised when a JWKS endpoint returns no key matching the requested ``kid``.""" +class IdJagIsNotABearerTokenError(Exception): + """Raised when bearer auth is handed an ID-JAG, which only the token endpoint may redeem.""" + + class JWKSUnreachableError(Exception): """Raised when an IdP's JWKS / OIDC discovery endpoint is unreachable and no cached copy is left to fall back on.""" @@ -114,9 +119,23 @@ _CACHED_VALUE: Final = TypeAdapter(object) _CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) +ID_JAG_TYP: Final = "oauth-id-jag+jwt" + + +@dataclass(frozen=True, slots=True) +class AssertionRejection: + """Why a JWT assertion failed verification. The reason is for the proxy log only, since + PyJWT's wording can quote the token's claims.""" + + reason: str + + class _JWTAuthSettings(Protocol): """The JWT auth settings block this handler reads back through ``getattr``, when one is configured.""" + @property + def id_jag_requesting_apps(self) -> Sequence[IdJagRequestingApp] | None: ... + @property def issuers(self) -> Sequence[JWTIssuerConfig] | None: ... @@ -282,7 +301,7 @@ class JWTHandler: return is_jwt_token(token=token) @staticmethod - def get_unverified_claims(token: str) -> dict | None: + def get_unverified_claims(token: str) -> dict[str, object] | None: """ Decode JWT claims without signature verification. Used for routing decisions before selecting validation path. @@ -949,7 +968,7 @@ class JWTHandler: keys: Final = response_json["keys"] if "keys" in response_json else response_json return cast(JWKKeyValue, keys) # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material - async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> dict: + async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> JWTKeyItem: resolved_jwks_url: Final = await self._resolve_jwks_url(jwks_url) keys: Final = await self._cached_with_stale_fallback( cache_key=f"litellm_jwt_auth_keys_{resolved_jwks_url}", @@ -960,11 +979,11 @@ class JWTHandler: public_key: Final = self.parse_keys(keys=keys, kid=kid) if public_key is not None: - return cast(dict, public_key) + return public_key raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={resolved_jwks_url}, kid={kid}") - async def get_public_key(self, kid: str | None) -> dict: + async def get_public_key(self, kid: str | None) -> JWTKeyItem: keys_url: Final = os.getenv("JWT_PUBLIC_KEY_URL") if keys_url is None: @@ -1151,7 +1170,7 @@ class JWTHandler: # its jwks_uri, matching JWTIssuerConfig.jwks_url's documented fallback. return f"{issuer_config.issuer.rstrip('/')}/.well-known/openid-configuration" - def _get_claim_value_for_issuer_mapping(self, token: dict, claim_field: str) -> object: + def _get_claim_value_for_issuer_mapping(self, token: Mapping[str, object], claim_field: str) -> object: """Resolve a mapped claim from ``token``. Returns ``None`` when the field is absent or empty so that mapped claims @@ -1168,8 +1187,10 @@ class JWTHandler: return None return claim_value - def _apply_issuer_claim_mappings(self, token: dict, issuer_config: JWTIssuerConfig) -> dict: - normalized: Final[dict] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS} + def _apply_issuer_claim_mappings( + self, token: Mapping[str, object], issuer_config: JWTIssuerConfig + ) -> dict[str, object]: + normalized: Final[dict[str, object]] = {k: v for k, v in token.items() if k not in self.LITELLM_INTERNAL_CLAIMS} normalized[self.LITELLM_JWT_ISSUER_CLAIM] = issuer_config.issuer claim_mappings: Final = [ (issuer_config.user_id_jwt_field, self.LITELLM_USER_ID_CLAIM), @@ -1192,8 +1213,8 @@ class JWTHandler: return normalized - def _get_jwk_from_public_key(self, public_key: dict) -> dict: - jwk: Final = {} + def _get_jwk_from_public_key(self, public_key: Mapping[str, object]) -> dict[str, object]: + jwk: Final[dict[str, object]] = {} for key in ["kty", "kid", "n", "e", "x", "y", "crv"]: if key in public_key: jwk[key] = public_key[key] @@ -1223,12 +1244,12 @@ class JWTHandler: def _decode_jwt_with_public_key( self, token: str, - public_key: dict | str, + public_key: JWTKeyItem | str, audience: str | list[str] | None, issuer: str | None = None, options: dict | None = None, disable_audience_validation: bool = False, - ) -> dict: + ) -> dict[str, object]: decode_options: Final = ( options if options is not None @@ -1298,11 +1319,64 @@ class JWTHandler: issuer_config=issuer_config, ) + async def decode_trusted_assertion( + self, token: str, audience: Sequence[str], required_claims: Sequence[str] + ) -> Mapping[str, object] | AssertionRejection: + """Verify a JWT assertion handed to the gateway's token endpoint with the issuer plumbing + JWT auth already trusts, a matching ``issuers`` entry or ``JWT_ISSUER`` with + ``JWT_PUBLIC_KEY_URL``, but against the caller's audience instead of the API audience. + An assertion whose ``iss`` nothing trusts is refused rather than checked unscoped. A JWKS + outage propagates as the 503 ``ProxyException`` so the caller does not call the token bad.""" + try: + header: Final = jwt.get_unverified_header(token) + except jwt.PyJWTError as malformed: + return AssertionRejection(reason=f"malformed assertion: {malformed}") + claims: Final = self.get_unverified_claims(token=token) + issuer_config: Final = self._get_configured_issuer(token=token) + trusted_issuer: Final = issuer_config.issuer if issuer_config is not None else os.getenv("JWT_ISSUER") + if not trusted_issuer or claims is None or claims.get("iss") != trusted_issuer: + return AssertionRejection(reason="iss is not an issuer this gateway trusts") + kid: Final = header.get("kid") + try: + public_key: Final = ( + await self._get_public_key_from_jwks_url(self._get_jwks_url_for_issuer(issuer_config), kid) + if issuer_config is not None + else await self.get_public_key(kid=kid) + ) + except JWKSUnreachableError as e: + raise jwks_unavailable_exception(e) from e + except Exception as no_key: # noqa: BLE001 # the key lookups raise plain Exceptions for a missing key + return AssertionRejection(reason=f"no signing key for the assertion: {no_key}") + try: + payload: Final = self._decode_jwt_with_public_key( + token=token, + public_key=public_key, + audience=list(audience), + issuer=trusted_issuer, + options={"require": list(required_claims)}, + ) + except jwt.PyJWTError as invalid: + return AssertionRejection(reason=f"assertion failed verification: {invalid}") + verified: Final = ( + payload + if issuer_config is None + else self._apply_issuer_claim_mappings(token=payload, issuer_config=issuer_config) + ) + validate: Final = self.litellm_jwtauth.custom_validate + if validate is not None and not validate(dict(verified)): + return AssertionRejection(reason="custom_validate rejected the assertion's claims") + return verified + async def auth_jwt(self, token: str) -> dict: header: Final = jwt.get_unverified_header(token) verbose_proxy_logger.debug("header: %s", header) + if header.get("typ") == ID_JAG_TYP: + raise IdJagIsNotABearerTokenError( + "an ID-JAG is an authorization grant, not a bearer token; redeem it at the token endpoint" + ) + kid: Final = header.get("kid", None) issuer_config: Final = self._get_configured_issuer(token=token) 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/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index ae129ca47b5..8dab693cd58 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -69,7 +69,8 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( # paths directly because the request route carries the resolved key id. _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spend") -_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", "llm_api_routes")) +_LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) +_AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) # Method-scoped because the generic credential CRUD handlers share these paths: # the route check cannot admit PATCH/DELETE on a name like "user_connections" @@ -154,7 +155,7 @@ class RouteChecks: ################################################ # For llm_api_routes, also check registered pass-through endpoints ################################################ - if allowed_route == "llm_api_routes": + if allowed_route in _LLM_API_ROUTE_GROUPS: if route == "/auto_router/session" and RouteChecks.get_request_method(request) == "GET": return True @@ -185,7 +186,9 @@ class RouteChecks: # headers, env, credentials). POST/PUT/DELETE on # these paths are admin-only management writes and # are intentionally not covered. - if RouteChecks._is_get_mcp_server_discovery_route(route=route, request=request): + if allowed_route == "llm_api_routes" and RouteChecks._is_get_mcp_server_discovery_route( + route=route, request=request + ): return True # Agent registry CRUD moved from llm_api_routes into @@ -195,7 +198,7 @@ class RouteChecks: # could reach these paths, so keep them reachable here; # the handlers in agent_endpoints/endpoints.py still # enforce proxy-admin on writes and scope reads by role. - if RouteChecks.check_route_access( + if allowed_route == "llm_api_routes" and RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.agent_management_routes.value, ): 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..61c26cbd502 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 @@ -504,6 +506,10 @@ async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) pass +def _encodes_as_latin1(value: str) -> bool: + return all(ord(character) < 256 for character in value) + + @lru_cache(maxsize=512) def _litellm_model_supports_stream_options(litellm_model: str) -> bool: try: @@ -633,6 +639,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) @@ -1968,10 +2031,11 @@ class ProxyBaseLLMRequestProcessing: headers.update(logging_caching_headers) try: - return {key: str(value) for key, value in headers.items() if value not in exclude_values} + present: Final = {key: str(value) for key, value in headers.items() if value not in exclude_values} except Exception as e: verbose_proxy_logger.error("Error setting custom headers: %s", e) return {} + return {key: value for key, value in present.items() if _encodes_as_latin1(value)} @staticmethod async def build_litellm_proxy_success_headers_from_llm_response( 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/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 81338988c66..7ac414bec00 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -18,9 +18,10 @@ import random from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from itertools import groupby -from typing import TYPE_CHECKING, Any, Final, Protocol +from itertools import chain, groupby +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import LiteralString from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS @@ -47,21 +48,73 @@ class ToolUsageTransaction: def response_tool_call_names(completion_response: object) -> tuple[str, ...]: """Tool names invoked in a completion response, in call order, for any response surface get_tool_calls_from_response understands (chat completions, Responses - API output items, Anthropic Messages tool_use blocks). Reads every choice of - an ``n>1`` chat response: each choice cost money and its tool calls ran.""" + API output items, Anthropic Messages tool_use blocks), plus the MCP tools an + auto-executed Responses API call ran. Reads every choice of an ``n>1`` chat + response: each choice cost money and its tool calls ran.""" if completion_response is None or isinstance(completion_response, Exception): return () from litellm.litellm_core_utils.prompt_templates.factory import ( get_tool_calls_from_response, ) + called_names: Final = ( + tool_call.get("name") + for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) + ) return tuple( stripped - for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) - if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) + for name in chain(called_names, _auto_executed_mcp_tool_names(completion_response)) + if isinstance(name, str) and (stripped := name.strip()) ) +class _ExecutedTool(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + + +class _OutputText(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + text: str + + +class _ToolExecutionResults(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: Literal["tool_execution_results"] + content: tuple[_OutputText, ...] + + +class _ResponsesApiOutput(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + output: tuple[object, ...] + + +_EXECUTED_TOOLS: Final = TypeAdapter(tuple[_ExecutedTool, ...]) + + +def _auto_executed_mcp_tool_names(completion_response: object) -> tuple[str | None, ...]: + try: + response: Final = _ResponsesApiOutput.model_validate(completion_response) + except ValidationError: + return () + return tuple(chain.from_iterable(_executed_tool_names(item) for item in response.output)) + + +def _executed_tool_names(output_item: object) -> tuple[str | None, ...]: + try: + results: Final = _ToolExecutionResults.model_validate(output_item) + executed: Final = tuple( + chain.from_iterable(_EXECUTED_TOOLS.validate_json(content.text) for content in results.content) + ) + except ValidationError: + return () + return tuple(tool.name for tool in executed) + + def build_tool_usage_transaction( request_id: str, start_time_iso: str, 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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 675bf0830f9..3ad5caade1b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -945,7 +945,9 @@ def apply_missing_session_id_policy( ) -def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: +def should_auto_drop_params_for_agentic_cli( + user_agent: str, caller_set_drop_params: bool, proxy_config: ProxyConfig +) -> bool: """drop_params defaults to on for agentic CLIs so their client-specific params (e.g. Claude Code's thinking, Codex's service_tier) don't fail requests routed to providers that reject them. An explicit drop_params @@ -955,7 +957,7 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False - if "drop_params" in data: + if caller_set_drop_params: return False config: Final = getattr(proxy_config, "config", None) litellm_settings: Final = config.get("litellm_settings") if isinstance(config, dict) else None @@ -2617,7 +2619,7 @@ async def add_litellm_data_to_request( user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent - if should_auto_drop_params_for_agentic_cli(user_agent, data, proxy_config): + if should_auto_drop_params_for_agentic_cli(user_agent, "drop_params" in data, proxy_config): data["drop_params"] = True # Merge caller-supplied tags (x-litellm-tags header, data["tags"] root-level) 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/sso/id_jag_assertion_capture.py b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py index 404fdfc83a9..ff7e68c3ab2 100644 --- a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py +++ b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py @@ -1,25 +1,26 @@ """Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion. An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693 -subject token. Only the generic OIDC login path reaches a token response the gateway retains one -from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion -and every store-sourced ID-JAG exchange fails for every user, however many times they sign in. -Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows -nothing about MCP. This module is the one shared answer both warn from. +subject token. The Google, Microsoft and generic OIDC login paths all reach a token response the +gateway retains one from, and the SAML login retains the signed assertion itself as a ``saml2`` +subject token. A deployment with no SSO never signs anyone in, so every store-sourced ID-JAG +exchange fails for every user. Neither side can see that alone: the MCP registration knows nothing +about SSO and the login knows nothing about MCP. This module is the one shared answer both warn from. """ from __future__ import annotations import os from enum import Enum +from typing import Final from typing_extensions import assert_never from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler -_GENERIC_OIDC_REMEDY = ( - "Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token " - "response the gateway retains an id_token from" +_OIDC_REMEDY: Final = ( + "Point SSO at an OIDC provider (GOOGLE_CLIENT_ID, MICROSOFT_CLIENT_ID or GENERIC_CLIENT_ID) or at SAML, " + "the login paths the gateway retains an identity assertion from" ) @@ -55,18 +56,14 @@ def id_jag_assertion_capture_gap() -> str | None: or ``None`` when that provider does capture an assertion.""" provider = active_sso_provider() match provider: - case ActiveSSOProvider.generic: + case ( + ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.generic | ActiveSSOProvider.saml + ): return None case ActiveSSOProvider.none: return ( "no SSO provider is configured, so no IdP identity assertion is ever captured and " - f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}" - ) - case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml: - return ( - f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no " - "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter " - f"how often they sign in. {_GENERIC_OIDC_REMEDY}" + f"ID-JAG credential resolution fails for every user. {_OIDC_REMEDY}" ) case _: assert_never(provider) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 4991436edac..ba265acfe4a 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -27,6 +27,7 @@ import hashlib import os import secrets import time +from collections.abc import Mapping from typing import Final, cast from urllib.parse import parse_qsl @@ -37,6 +38,10 @@ from pydantic import ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SSOIdentityAssertion, + assertion_from_saml_login, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.management_endpoints.sso_helper_utils import SSO_SESSIONS_TARGET from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role @@ -212,7 +217,7 @@ class SAMLAuthHandler: return OneLogin_Saml2_IdPMetadataParser.merge_settings(sp_settings, idp_settings) @staticmethod - def _prepare_request_data(request: Request, post_data: dict[str, str] | None = None) -> dict[str, object]: + def _prepare_request_data(request: Request, post_data: Mapping[str, str] | None = None) -> dict[str, object]: base: Final = SAMLAuthHandler._base_url(request) scheme, _, host_part = base.partition("://") host: Final = host_part.split("/", 1)[0] @@ -228,7 +233,7 @@ class SAMLAuthHandler: async def _build_auth( request: Request, cache: DualCache, - post_data: dict[str, str] | None = None, + post_data: Mapping[str, str] | None = None, ) -> "OneLogin_Saml2_Auth": if not SAML_AVAILABLE: raise _saml_unavailable_error() @@ -310,7 +315,9 @@ class SAMLAuthHandler: return dict(parse_qsl(body.decode("utf-8", "replace"))) @staticmethod - async def handle_acs(request: Request, cache: DualCache, post_data: dict[str, str]) -> CustomOpenID: + async def handle_acs( + request: Request, cache: DualCache, post_data: Mapping[str, str] + ) -> tuple[CustomOpenID, SSOIdentityAssertion | None]: auth: Final = await SAMLAuthHandler._build_auth(request, cache, post_data=post_data) browser_request_id: Final = request.cookies.get(_SAML_AUTHN_STATE_COOKIE) try: @@ -330,7 +337,50 @@ class SAMLAuthHandler: ) await SAMLAuthHandler._enforce_response_binding(auth, cache, browser_request_id) - return SAMLAuthHandler._result_from_auth(auth) + return SAMLAuthHandler._result_from_auth(auth), SAMLAuthHandler._retained_assertion(auth) + + @staticmethod + def _retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None: + """The one signed ```` the validated response carried, serialized on its own + so its enveloped signature still verifies at the IdP's token endpoint. Retention is optional + egress material, so a failure here is logged and the validated login goes on without it.""" + try: + return SAMLAuthHandler._parse_retained_assertion(auth) + except Exception as e: # noqa: BLE001 - lxml and python3-saml expose no common base; capture must not block login + verbose_proxy_logger.warning( + "SAML login validated but its assertion could not be retained for oauth2_id_jag MCP servers: %s", e + ) + return None + + @staticmethod + def _parse_retained_assertion(auth: "OneLogin_Saml2_Auth") -> SSOIdentityAssertion | None: + xml: Final = cast(bytes | None, auth.get_last_response_xml()) # cast-ok: untyped python3-saml + if not xml: + verbose_proxy_logger.warning( + "SAML login validated but the response XML was not retained, so this user has no identity " + "assertion for oauth2_id_jag MCP servers." + ) + return None + root: Final = OneLogin_Saml2_XML.to_etree(xml) + assertions: Final = OneLogin_Saml2_XML.query(root, "//saml:Assertion") + if len(assertions) != 1: + encrypted: Final = OneLogin_Saml2_XML.query(root, "//saml:EncryptedAssertion") + verbose_proxy_logger.warning( + "SAML login validated but the response carried %d plaintext saml:Assertion elements and %d " + "saml:EncryptedAssertion elements, so no identity assertion was retained for oauth2_id_jag MCP " + "servers. Those servers need exactly one signed, unencrypted assertion from the IdP.", + len(assertions), + len(encrypted), + ) + return None + issuers: Final = OneLogin_Saml2_XML.query(assertions[0], "./saml:Issuer") + issuer: Final = cast(str | None, issuers[0].text) if issuers else None # cast-ok: untyped python3-saml + not_on_or_after: Final = auth.get_last_assertion_not_on_or_after() + return assertion_from_saml_login( + assertion_xml=cast(bytes, OneLogin_Saml2_XML.to_string(assertions[0], encoding="utf-8")), # cast-ok: saml + not_on_or_after=not_on_or_after if isinstance(not_on_or_after, int) else None, + issuer=issuer, + ) @staticmethod def _replay_guard_ttl(auth: "OneLogin_Saml2_Auth") -> int: 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 4629103070f..4cc849762aa 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -37,10 +37,12 @@ from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: import httpx + from litellm.caching.redis_cache import RedisCache + import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse -from pydantic import ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator import litellm from litellm._internal_context import with_service_target @@ -69,7 +71,12 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import ( + exchange_login_assertion_if_configured, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + GOOGLE_REFRESH_TOKEN_PARAMS, + MICROSOFT_REFRESH_TOKEN_SCOPE, SSOIdentityAssertion, assertion_from_sso_login, ema_assertion_retention_enabled, @@ -82,6 +89,7 @@ from litellm.proxy._types import ( Member, NewTeamRequest, NewUserRequest, + NewUserRequestTeam, NewUserResponse, ProxyErrorTypes, ProxyException, @@ -111,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, @@ -1786,6 +1798,19 @@ async def warn_if_id_jag_assertion_uncaptured( ) +async def login_asks_for_refresh_token(retention_enabled: RetentionCheck | None = None) -> bool: + """A refresh token is only issued when the authorization request asks for one, and asking costs the + user an offline-access consent, so the Google and Microsoft logins ask only while an ``oauth2_id_jag`` + server exists whose assertion the refresher would renew. A store outage answers no: a login must not + fail over a renewal that is only an optimisation on top of signing in again.""" + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + return await check() + except Exception as exc: # noqa: BLE001 # the login must not fail over the renewal gate + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers before SSO login: %s", exc) + return False + + async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None = None) -> None: gap: Final = id_jag_assertion_capture_gap() if gap is None: @@ -2125,6 +2150,50 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role +async def _sso_provider_callback( + request: Request, + redirect_url: str, + *, + google_client_id: str | None, + microsoft_client_id: str | None, + generic_client_id: str | None, + jwt_handler: JWTHandler, + sso_jwt_handler: JWTHandler | None, + return_raw_sso_response: bool = False, +) -> tuple[ + OpenID | dict[str, object] | None, dict[str, object] | None, dict[str, object] | None, SSOIdentityAssertion | None +]: + """Dispatch the callback to the configured provider in the order the login redirect used (Google, + Microsoft, generic) and hand back what every branch has to agree on: the parsed identity, the raw + userinfo and access-token claims only the generic branch exposes, and the IdP identity assertion + an ``oauth2_id_jag`` MCP server spends. ``None`` identity means no provider is configured.""" + if google_client_id is not None: + google_result, google_assertion = await GoogleSSOHandler.get_google_callback_response( + request=request, + google_client_id=google_client_id, + redirect_url=redirect_url, + return_raw_sso_response=return_raw_sso_response, + ) + return google_result, None, None, google_assertion + if microsoft_client_id is not None: + microsoft_result, microsoft_assertion = await MicrosoftSSOHandler.get_microsoft_callback_response( + request=request, + microsoft_client_id=microsoft_client_id, + redirect_url=redirect_url, + return_raw_sso_response=return_raw_sso_response, + ) + return microsoft_result, None, None, microsoft_assertion + if generic_client_id is not None: + return await get_generic_sso_response( + request=request, + jwt_handler=jwt_handler, + generic_client_id=generic_client_id, + redirect_url=redirect_url, + sso_jwt_handler=sso_jwt_handler, + ) + return None, None, None, None + + @router.get("/sso/callback", tags=["experimental"], include_in_schema=False) async def auth_callback(request: Request, state: str | None = None): """Verify login""" @@ -2173,9 +2242,6 @@ async def auth_callback(request: Request, state: str | None = None): microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None) google_client_id: Final = os.getenv("GOOGLE_CLIENT_ID", None) generic_client_id: Final = os.getenv("GENERIC_CLIENT_ID", None) - received_response: dict | None = None - access_token_payload: dict | None = None - sso_assertion: SSOIdentityAssertion | None = None # get url from request if master_key is None: raise ProxyException( @@ -2187,33 +2253,15 @@ async def auth_callback(request: Request, state: str | None = None): redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(request=request, sso_callback_route="sso/callback") verbose_proxy_logger.info("Redirecting to %s", redirect_url) - result = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - sso_assertion, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) + result, received_response, access_token_payload, sso_assertion = await _sso_provider_callback( + request, + redirect_url, + google_client_id=google_client_id, + microsoft_client_id=microsoft_client_id, + generic_client_id=generic_client_id, + jwt_handler=jwt_handler, + sso_jwt_handler=sso_jwt_handler, + ) if result is None: raise HTTPException( @@ -2303,7 +2351,7 @@ async def saml_callback(request: Request): if "SAMLResponse" not in post_data: raise HTTPException(status_code=400, detail="Missing SAMLResponse in callback request.") - result: Final = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) + user, retained = await SAMLAuthHandler.handle_acs(request=request, cache=saml_sso_cache, post_data=post_data) await _raise_if_sso_exceeds_free_user_limit(premium_user, prisma_client) @@ -2316,7 +2364,7 @@ async def saml_callback(request: Request): ) return await SSOAuthenticationHandler.get_redirect_response_from_openid( - result=result, + result=user, request=request, received_response=None, generic_client_id=None, @@ -2324,6 +2372,7 @@ async def saml_callback(request: Request): access_token_payload=None, jwt_handler=jwt_handler, return_to=cp_return_to, + sso_assertion=retained, ) @@ -2440,7 +2489,8 @@ async def _complete_cli_sso_callback_session( raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") resolved_user_id: Final = _require_sso_user_id(user_info.user_id) - await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=sso_assertion) + retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion) + await retain_sso_identity_assertion_for_ema(user_id=resolved_user_id, assertion=retained_assertion) await warn_if_id_jag_assertion_uncaptured(sso_assertion) teams: list[str] = [] @@ -2780,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"], @@ -2915,8 +3011,8 @@ def _is_same_origin_return_path(return_to: str) -> bool: async def sso_return_to_redirect( return_to: str | None, jwt_token: str, - redis_usage_cache, - user_api_key_cache, + redis_usage_cache: "RedisCache | None", + user_api_key_cache: UserApiKeyCache, request: Request, ) -> RedirectResponse | None: """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard. @@ -3063,6 +3159,7 @@ class SSOAuthenticationHandler: generic_client_id: str | None = None, state: str | None = None, request: Request | None = None, + retention_enabled: RetentionCheck | None = None, ) -> RedirectResponse | None: """ Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url` @@ -3074,10 +3171,13 @@ class SSOAuthenticationHandler: generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None. request: Optional FastAPI request, used to drive the ``Secure`` attribute on the ``litellm_oauth_state`` CSRF cookie. + retention_enabled: Whether an ``oauth2_id_jag`` server exists, deciding if the Google and + Microsoft logins ask the IdP for a refresh token. Defaults to the store's own check. Returns: RedirectResponse: The redirect response from the SSO provider. """ + asks_for_refresh_token: Final = await login_asks_for_refresh_token(retention_enabled) # Google SSO Auth if google_client_id is not None: google_client_secret: Final = os.getenv("GOOGLE_CLIENT_SECRET", None) @@ -3099,7 +3199,10 @@ class SSOAuthenticationHandler: google_client_id, ) with google_sso: - return await google_sso.get_login_redirect(state=state) + return await google_sso.get_login_redirect( + state=state, + params=dict(GOOGLE_REFRESH_TOKEN_PARAMS) if asks_for_refresh_token else None, + ) # Microsoft SSO Auth elif microsoft_client_id is not None: microsoft_client_secret: Final = os.getenv("MICROSOFT_CLIENT_SECRET", None) @@ -3117,6 +3220,7 @@ class SSOAuthenticationHandler: tenant=microsoft_tenant, redirect_uri=redirect_url, allow_insecure_http=True, + scope=[*CustomMicrosoftSSO.scope, MICROSOFT_REFRESH_TOKEN_SCOPE] if asks_for_refresh_token else None, ) with microsoft_sso: return await microsoft_sso.get_login_redirect(state=state) @@ -3418,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: @@ -3815,7 +3920,8 @@ class SSOAuthenticationHandler: ) if isinstance(user_id, str) and user_id: - await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=sso_assertion) + retained_assertion: Final = await exchange_login_assertion_if_configured(sso_assertion) + await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=retained_assertion) await warn_if_id_jag_assertion_uncaptured(sso_assertion) disabled_non_admin_personal_key_creation: Final = get_disabled_non_admin_personal_key_creation() @@ -4399,6 +4505,21 @@ class SSOAuthenticationHandler: return userinfo +class _SSOTokenResponse(BaseModel): + """The id_token and refresh token fastapi-sso keeps on the SSO instance after the code exchange, + read through a validated view because the package ships no type information.""" + + model_config = ConfigDict(frozen=True, from_attributes=True) + + id_token: str | None = None + refresh_token: str | None = None + + @field_validator("id_token", "refresh_token", mode="before") + @classmethod + def _strings_only(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + + class MicrosoftSSOHandler: """ Handles Microsoft SSO callback response and returns a CustomOpenID object @@ -4433,9 +4554,10 @@ class MicrosoftSSOHandler: microsoft_client_id: str, redirect_url: str, return_raw_sso_response: bool = False, - ) -> CustomOpenID | OpenID | dict: + ) -> tuple[CustomOpenID | OpenID | dict[str, object], SSOIdentityAssertion | None]: """ - Get the Microsoft SSO callback response + Get the Microsoft SSO callback response, with the IdP identity assertion the token + response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends) Args: return_raw_sso_response: If True, return the raw SSO response @@ -4470,6 +4592,8 @@ class MicrosoftSSOHandler: ) or {} ) + microsoft_tokens: Final = _SSOTokenResponse.model_validate(microsoft_sso) + sso_assertion: Final = assertion_from_sso_login(microsoft_tokens.id_token, microsoft_tokens.refresh_token) if not return_raw_sso_response and "error" in original_msft_result and original_msft_result["error"]: verbose_proxy_logger.warning("Microsoft SSO login failed: Graph /me returned an error instead of a profile") raise HTTPException( @@ -4494,7 +4618,7 @@ class MicrosoftSSOHandler: if return_raw_sso_response: original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = user_team_ids original_msft_result["app_roles"] = app_roles - return original_msft_result or {} + return original_msft_result or {}, sso_assertion from litellm.proxy.management_endpoints.sso.agent_subject_enrollment import microsoft_interactive_subject @@ -4517,7 +4641,7 @@ class MicrosoftSSOHandler: team_ids=user_team_ids, user_role=user_role, ) - return result + return result, sso_assertion @staticmethod def openid_from_response( @@ -4793,9 +4917,10 @@ class GoogleSSOHandler: google_client_id: str, redirect_url: str, return_raw_sso_response: bool = False, - ) -> OpenID | dict: + ) -> tuple[OpenID | dict[str, object], SSOIdentityAssertion | None]: """ - Get the Google SSO callback response + Get the Google SSO callback response, with the IdP identity assertion the token + response carried (the id_token and refresh token an ``oauth2_id_jag`` server spends) Args: return_raw_sso_response: If True, return the raw SSO response @@ -4816,17 +4941,13 @@ class GoogleSSOHandler: ) # if user is trying to get the raw sso response for debugging, return the raw sso response - if return_raw_sso_response: - return ( - await google_sso.verify_and_process( - request=request, - convert_response=False, - ) - or {} - ) - - result: Final = await google_sso.verify_and_process(request) - return result or {} + result: Final = ( + await google_sso.verify_and_process(request=request, convert_response=False) + if return_raw_sso_response + else await google_sso.verify_and_process(request) + ) + google_tokens: Final = _SSOTokenResponse.model_validate(google_sso) + return result or {}, assertion_from_sso_login(google_tokens.id_token, google_tokens.refresh_token) def _raise_if_sso_debug_disabled() -> None: @@ -4928,37 +5049,16 @@ async def debug_sso_callback(request: Request): else: redirect_url += "/sso/debug/callback" - result = None - received_response: dict | None = None - access_token_payload: dict | None = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - _sso_assertion, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) + result, received_response, access_token_payload, _sso_assertion = await _sso_provider_callback( + request, + redirect_url, + google_client_id=google_client_id, + microsoft_client_id=microsoft_client_id, + generic_client_id=generic_client_id, + jwt_handler=jwt_handler, + sso_jwt_handler=sso_jwt_handler, + return_raw_sso_response=True, + ) # If result is None, return a basic error message if result is None: 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/proxy_server.py b/litellm/proxy/proxy_server.py index 9eb2bc93391..3de3dccc532 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -171,6 +171,7 @@ from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, + CountContextManagement, ModelResponse, ModelResponseStream, StreamingChoices, @@ -10012,6 +10013,7 @@ def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]: _STREAM_KEEPALIVE: Final = object() _KEEPALIVE_MIN_SECONDS: Final = 1.0 +_STREAM_LOOP_HOLD_SECONDS: Final = 0.05 _KEEPALIVE_MAX_SECONDS: Final = 300.0 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) _EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) @@ -10287,7 +10289,14 @@ async def async_data_generator( else stream_iterator ) + loop_released_at = time.monotonic() # rebind-ok: last time this stream let other tasks run async for item in stream_source: + # Unblock the event loop: a stream whose chunks are already buffered never awaits + # between chunks, so without this yield one long stream holds the loop for its whole body. + # A real await in between may already have released the loop; this only guarantees it + if time.monotonic() - loop_released_at >= _STREAM_LOOP_HOLD_SECONDS: + await asyncio.sleep(0) + loop_released_at = time.monotonic() if item is _STREAM_KEEPALIVE: yield ": ping\n\n" continue @@ -14456,6 +14465,7 @@ async def _try_provider_token_count( system: str | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14471,6 +14481,7 @@ async def _try_provider_token_count( system=system, messages_format=messages_format, output_format=output_format, + context_management=context_management, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -14501,6 +14512,39 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) +def _router_default_litellm_params(router: Router) -> Mapping[str, object]: + defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None) + if defaults is None: + return {} + return {key: value for key, value in defaults.items() if value is not None} + + +def _count_context_management( + context_management: JsonValue, model: str, request_litellm_params: Mapping[str, JsonValue] +) -> CountContextManagement | None: + from litellm.llms.anthropic.pass_through.adapters.handler import context_management_explicitly_dropped + + if context_management is None or llm_router is None: + return None + inference_model: Final = resolve_inference_model( + model, general_settings_view(), user_model if isinstance(user_model, str) else None + ) + router_defaults: Final = _router_default_litellm_params(llm_router) + sent_params: Final = tuple( + {**deployment["litellm_params"], **router_defaults, **request_litellm_params} + for deployment in llm_router.get_routable_deployments( + inference_model if isinstance(inference_model, str) else model + ) + ) + if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params): + return None + return CountContextManagement( + spec=context_management, + drop_invalid_edits=litellm.drop_params is True + or any(params.get("drop_params") is True for params in sent_params), + ) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14531,6 +14575,8 @@ async def count_request_tokens( call_endpoint: bool = False, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: JsonValue = None, + request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -14602,6 +14648,7 @@ async def count_request_tokens( system=system, messages_format=messages_format, output_format=output_format, + context_management=_count_context_management(context_management, request.model, request_litellm_params), ) if result is not None: return result @@ -19698,6 +19745,26 @@ async def delete_config_general_settings( return response +def _configured_callback_names(value: object) -> tuple[JsonValue, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(cast(Iterable[JsonValue], value)) # cast-ok: config.yaml values are untyped + return () + + +def _is_callback_name(entry: object, callback_name: str) -> bool: + return isinstance(entry, str) and entry.lower() == callback_name + + +def _callback_lists_to_rewrite( + matching_keys: tuple[str, ...], owned_by_config: Callable[[str], bool] +) -> tuple[str, ...]: + """The matching lists the database owns, or all of them when the config file owns every one.""" + database_owned: Final = tuple(key for key in matching_keys if not owned_by_config(key)) + return database_owned or matching_keys + + @router.post( "/config/callback/delete", tags=["config.yaml"], @@ -19733,24 +19800,42 @@ async def delete_callback( try: # Get current configuration - config: Final = await proxy_config.get_config() + config: Final = cast( # cast-ok: get_config returns an untyped config dict + dict[str, object], + await proxy_config.get_config(), + ) callback_name: Final = data.callback_name.lower() # Check if callback exists in current configuration - litellm_settings: Final = config.get("litellm_settings", {}) - success_callbacks: Final = litellm_settings.get("success_callback", []) + litellm_settings: Final[Mapping[str, object]] = cast( # cast-ok: the config dict's values are untyped + Mapping[str, object], + config.get("litellm_settings", {}), + ) + configured_lists: Final = { + key: _configured_callback_names(litellm_settings.get(key)) for key in _DB_CONFIG_CALLBACK_EVENT_TYPES + } + matching_keys: Final = tuple( + key + for key, names in configured_lists.items() + if any(_is_callback_name(entry, callback_name) for entry in names) + ) - if callback_name not in success_callbacks: + if not matching_keys: raise HTTPException( status_code=404, detail={"error": f"Callback '{callback_name}' not found in active configuration"}, ) - before_success_callbacks: Final = list(success_callbacks) - - # Remove callback from success_callback list - success_callbacks.remove(callback_name) - config.setdefault("litellm_settings", {})["success_callback"] = success_callbacks + rewritten_keys: Final = _callback_lists_to_rewrite(matching_keys, proxy_config.litellm_settings.owned_by_config) + remaining_lists: Final = MappingProxyType( + { + key: tuple(entry for entry in configured_lists[key] if not _is_callback_name(entry, callback_name)) + for key in rewritten_keys + } + ) + before_callbacks: Final[JsonValue] = {key: list(configured_lists[key]) for key in rewritten_keys} + after_callbacks: Final[JsonValue] = {key: list(names) for key, names in remaining_lists.items()} + config["litellm_settings"] = {**litellm_settings, **after_callbacks} # Save the updated configuration await proxy_config.save_config(new_config=config) @@ -19759,8 +19844,8 @@ async def delete_callback( create_config_audit_log( "litellm_settings", "deleted", - {"success_callback": before_success_callbacks}, - {"success_callback": success_callbacks}, + before_callbacks, + after_callbacks, user_api_key_dict, ) ) @@ -19771,7 +19856,7 @@ async def delete_callback( return { "message": f"Successfully deleted callback: {callback_name}", "removed_callback": callback_name, - "remaining_callbacks": success_callbacks, + "remaining_callbacks": tuple(chain.from_iterable(remaining_lists.values())), "deleted_at": datetime.now().isoformat(), } diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4a6a88f9a86..cb9e193fa48 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2417,6 +2417,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "Nace", + "provider_display_name": "Nace", + "litellm_provider": "nace", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://console.nace.ai", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "nace/drex-latest" + }, { "provider": "NEBIUS", "provider_display_name": "Nebius", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index a44f1d10fda..207e966ef28 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 @@ -57,6 +63,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -69,7 +76,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) @@ -1376,6 +1382,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": { @@ -1432,6 +1447,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. @@ -1464,6 +1481,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, @@ -1471,9 +1495,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/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9ebea589cc9..c7ca08fb513 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -47,6 +47,7 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.hidden_params import served_from_cache from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call from litellm.litellm_core_utils.litellm_logging import ( coerce_model_access_groups, @@ -542,6 +543,7 @@ def get_logging_payload( kwargs = {} rejected_as_unknown_model: Final = isinstance(response_obj, ProxyModelNotFoundError) + response_served_from_cache: Final = served_from_cache(response_obj) if response_obj is None: response_obj = {} elif not isinstance(response_obj, BaseModel) and not isinstance(response_obj, dict): @@ -812,7 +814,7 @@ def get_logging_payload( cache_key = litellm_params["preset_cache_key"] else: cache_key = litellm.cache.get_cache_key(**kwargs) - if cache_hit is True: + if cache_hit is True or response_served_from_cache: import time id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index b327129793d..4ece88f757d 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -343,7 +343,7 @@ class UISettings(LiteLLMBaseModel): ) enable_chat_ui: bool = Field( - default=False, + default=True, description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index be6532ba36b..36ace3c1697 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -65,6 +65,7 @@ from litellm.constants import ( SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, ) +from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( bug_report_notice, should_report_bug, @@ -990,6 +991,24 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) +_FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[tuple[Usage, float], ...]) + + +def _billed_across_attempts( + failed_attempts: tuple[tuple[Usage, float], ...], recovered_usage: object, recovered_cost: object +) -> tuple[Usage, float]: + this_attempt: Final = ( + ((recovered_usage, float(recovered_cost) if isinstance(recovered_cost, (int, float)) else 0.0),) + if isinstance(recovered_usage, Usage) + else () + ) + attempts: Final = (*failed_attempts, *this_attempt) + return ( + BaseTokenUsageProcessor.combine_usage_objects([usage for usage, _ in attempts]), + sum(response_cost for _, response_cost in attempts), + ) + + def _failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], @@ -1004,6 +1023,9 @@ def _failure_usage_to_lift( chat-completions provider) without filling optional_params. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if failed_attempts: + return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL): @@ -3471,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 @@ -3535,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/rag/main.py b/litellm/rag/main.py index 5ab39bec080..be566a0da8b 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -11,8 +11,7 @@ __all__ = ["aingest", "aquery", "ingest", "query"] import asyncio import contextvars -from collections.abc import Coroutine, Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Coroutine, Mapping from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -21,7 +20,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm -from litellm._internal_context import is_internal_call +from litellm._internal_context import internal_sub_call from litellm.cost_calculator import vector_store_search_cost from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion @@ -212,25 +211,6 @@ async def aingest( ) -@contextmanager -def _suppressed_sub_call_billing() -> Iterator[None]: - """ - Suppress a sub-call's own billing event so the parent aquery event bills it. - - Every suppressed sub-call's cost must be folded into the parent event: - into the response's hidden response_cost on the non-streaming path, or via - the logging object's additional_response_cost on the streaming path (the - streamed cost is computed from assembled chunks after this pipeline - returns, so there is no response object to fold into here). - """ - previous: Final = is_internal_call.get() - is_internal_call.set(True) - try: - yield - finally: - is_internal_call.set(previous) - - async def _execute_query_pipeline( model: str, messages: list[AllMessageValues], @@ -272,7 +252,7 @@ async def _execute_query_pipeline( forwarded_search_params: Final = MappingProxyType( {**provider_search_params, **kwargs, **filter_search_params, **store_search_params} ) - with _suppressed_sub_call_billing(): + with internal_sub_call(): search_response: Final = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, @@ -306,7 +286,7 @@ async def _execute_query_pipeline( if rerank and rerank.get("enabled"): documents: Final = RAGQuery.extract_documents_from_search(search_response) if documents: - with _suppressed_sub_call_billing(): + with internal_sub_call(): rerank_response = await litellm.arerank( model=rerank["model"], query=query_text, @@ -324,7 +304,7 @@ async def _execute_query_pipeline( modified_messages: Final = messages[:-1] + [context_message] + [messages[-1]] # Use router if available to properly resolve virtual model names - with _suppressed_sub_call_billing(): + with internal_sub_call(): if router is not None: response = await router.acompletion( model=model, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 367d5be7bfc..96b4b081c05 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -37,12 +37,14 @@ from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ..litellm_core_utils.realtime_streaming import client_websocket_scope from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client -from ..llms.bedrock.realtime.handler import BedrockRealtime +from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed 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 +132,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 +193,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 +268,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 +435,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 + ), ) @@ -451,6 +472,7 @@ async def _arealtime( For PROXY use only. """ + raise_if_session_committed(client_websocket_scope(websocket)) headers = cast(dict | None, kwargs.get("headers")) extra_headers: Final = cast(dict | None, kwargs.get("extra_headers")) if headers is None: @@ -497,14 +519,17 @@ 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 + drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True provider_config: BaseRealtimeConfig | None = None if _custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_config( model=model, provider=LlmProviders(_custom_llm_provider), + drop_params=drop_params, ) if provider_config is not None: await base_llm_http_handler.async_realtime( @@ -556,8 +581,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, @@ -585,7 +611,6 @@ async def _arealtime( aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint") aws_external_id: Final = kwargs.get("aws_external_id") aws_session_tags: Final = kwargs.get("aws_session_tags") - drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True await bedrock_realtime.async_realtime( model=model, @@ -671,6 +696,7 @@ async def _arealtime( resolve_access_token=resolve_vertex_access_token, project=resolved_project, location=vertex_location, + drop_params=drop_params, ) await base_llm_http_handler.async_realtime( @@ -772,9 +798,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/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 0ab2052abb0..5489d2223ac 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -29,6 +29,9 @@ _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) class LiteLLMCompletionTransformationHandler: def response_api_handler( self, + previous_response_client_ids: tuple[str, ...] = (), + /, + *, model: str, input: str | ResponseInputParam, responses_api_request: ResponsesAPIOptionalRequestParams, @@ -61,6 +64,7 @@ class LiteLLMCompletionTransformationHandler: if _is_async: return self.async_response_api_handler( + previous_response_client_ids, litellm_completion_request=litellm_completion_request, request_input=bridged_input, responses_api_request=bridged_request, @@ -100,6 +104,9 @@ class LiteLLMCompletionTransformationHandler: async def async_response_api_handler( self, + previous_response_client_ids: tuple[str, ...] = (), + /, + *, litellm_completion_request: dict, request_input: str | ResponseInputParam, responses_api_request: ResponsesAPIOptionalRequestParams, @@ -112,6 +119,7 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_request=litellm_completion_request, instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + previous_response_client_ids=previous_response_client_ids, ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index e78ec5e5d3e..615bdb5bbfb 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -14,7 +14,7 @@ from litellm.types.llms.openai import ( GenericChatCompletionMessage, ResponseInputParam, ) -from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse +from litellm.types.utils import CallTypes, ChatCompletionMessageToolCall, Message, ModelResponse if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload @@ -42,6 +42,30 @@ def _normalize_redacted_tool_call_arguments(message: Message) -> None: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +_MCP_TOOL_CALL_TYPES: Final = frozenset((CallTypes.list_mcp_tools.value, CallTypes.call_mcp_tool.value)) + + +def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool: + return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES + + +_MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results")) + + +def _is_model_output(item: object) -> bool: + return not (isinstance(item, Mapping) and item.get("type") in _MCP_AUTO_EXECUTE_OUTPUT_TYPES) + + +def _stored_responses_api_output(stored_response: object) -> ResponseInputParam | None: + match stored_response: + case {"object": "response", "output": list() as output}: + return [item for item in output if _is_model_output(item)] + case {"object": "response"}: + return [] + case _: + return None + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -52,6 +76,7 @@ class ResponsesSessionHandler: async def get_chat_completion_message_history_for_previous_response_id( previous_response_id: str, item_cache_control: bool = False, + client_response_ids: tuple[str, ...] = (), ) -> ChatCompletionSession: """ Return the chat completion message history for a previous response id @@ -61,9 +86,13 @@ class ResponsesSessionHandler: ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") - all_spend_logs: list[ - SpendLogsPayload - ] = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(previous_response_id) + all_spend_logs: Final = [ + spend_log + for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id, client_response_ids + ) + if _is_a_conversation_turn(spend_log) + ] verbose_proxy_logger.debug("found %s spend logs for this response id", len(all_spend_logs)) litellm_session_id: str | None = None @@ -158,8 +187,17 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final = spend_log.get("response", "{}") - if isinstance(_response_output, dict) and _response_output and _response_output != {}: + _response_output: Final[object] = spend_log.get("response", "{}") + if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: + chat_completion_message_history.extend( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=responses_api_output, + responses_api_request={}, + replay_reasoning=True, + item_cache_control=item_cache_control, + ) + ) + elif isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) for choice in model_response.choices: @@ -176,7 +214,7 @@ class ResponsesSessionHandler: Get the parsed proxy server request from the spend log """ proxy_server_request: Final[str | dict] = spend_log.get("proxy_server_request") or "{}" - proxy_server_request_dict: dict | None = None + proxy_server_request_dict: dict[str, object] | None = None if isinstance(proxy_server_request, dict): proxy_server_request_dict = proxy_server_request else: @@ -229,7 +267,7 @@ class ResponsesSessionHandler: @staticmethod async def get_proxy_server_request_from_cold_storage_with_object_key( object_key: str, - ) -> dict | None: + ) -> dict[str, object] | None: """ Get the proxy server request from cold storage using the object key directly. @@ -251,7 +289,7 @@ class ResponsesSessionHandler: @staticmethod def _should_check_cold_storage_for_full_payload( - proxy_server_request_dict: dict | None, + proxy_server_request_dict: Mapping[str, object] | None, ) -> bool: """ Only check cold storage when both are true @@ -274,6 +312,7 @@ class ResponsesSessionHandler: @staticmethod async def get_all_spend_logs_for_previous_response_id( previous_response_id: str, + client_response_ids: tuple[str, ...] = (), ) -> "list[SpendLogsPayload]": """ Get all spend logs for a previous response id @@ -307,11 +346,12 @@ class ResponsesSessionHandler: if prisma_client is None: return [] + request_ids: Final = list(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) query: Final = """ WITH matching_session AS ( SELECT session_id, api_key, "user", team_id FROM "LiteLLM_SpendLogs" - WHERE request_id = $1 + WHERE request_id = ANY($1::text[]) ) SELECT * FROM "LiteLLM_SpendLogs" AS logs @@ -334,7 +374,7 @@ class ResponsesSessionHandler: for attempt in range(max_attempts): if attempt: await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL) - if spend_logs := await prisma_client.db.query_raw(query, response_id): + if spend_logs := await prisma_client.db.query_raw(query, request_ids): verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", response_id, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 418bdb3d6db..5c9e5cac5ab 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -100,6 +100,11 @@ def _delta_carries_output(delta: ChatCompletionDelta) -> bool: return any(fields) or _delta_has_signed_thinking_block(delta) +def _delta_carries_reasoning_and_answer(delta: ChatCompletionDelta) -> bool: + has_reasoning: Final = bool(getattr(delta, "reasoning_content", None)) or _delta_has_signed_thinking_block(delta) + return has_reasoning and bool(delta.content or delta.tool_calls) + + class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ Async iterator for processing streaming responses from the Responses API. @@ -566,6 +571,26 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._buffered_chunk = None return buffered + def _reasoning_part_first(self, chunk: ModelResponseStream) -> ModelResponseStream: + if not chunk.choices or not _delta_carries_reasoning_and_answer(chunk.choices[0].delta): + return chunk + first_choice: Final = chunk.choices[0] + delta: Final = first_choice.delta + answer_delta: Final = delta.model_copy(update={"reasoning_content": None, "thinking_blocks": None}) + self._buffered_chunk = chunk.model_copy( + update={"choices": [first_choice.model_copy(update={"delta": answer_delta}), *chunk.choices[1:]]} + ) + reasoning_delta: Final = ChatCompletionDelta( + reasoning_content=getattr(delta, "reasoning_content", None), + thinking_blocks=delta.thinking_blocks if hasattr(delta, "thinking_blocks") else None, + ) + return chunk.model_copy( + update={ + "provider_specific_fields": None, + "choices": [first_choice.model_copy(update={"delta": reasoning_delta, "finish_reason": None})], + } + ) + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -1052,6 +1077,41 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_message_item_added_events() return + def _queue_reasoning_done_events_on_reasoning_end(self, chunk: ModelResponseStream) -> None: + if not self._reasoning_active or self._reasoning_done_emitted: + return + delta: Final = chunk.choices[0].delta if chunk.choices else None + if delta and getattr(delta, "reasoning_content", None): + self._accumulated_reasoning_content_parts.append(delta.reasoning_content) + if not self._is_reasoning_end(chunk): + return + reasoning_content: Final = "".join(self._accumulated_reasoning_content_parts) + self._cached_reasoning_item_id = ( + self._reasoning_item_id or self._cached_reasoning_item_id or mint_reasoning_item_id() + ) + reasoning_item_id: Final = self._cached_reasoning_item_id + self._sequence_number += 1 + text_done_event: Final = self.create_reasoning_summary_text_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._sequence_number += 1 + part_done_event: Final = self.create_reasoning_summary_part_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._sequence_number += 1 + reasoning_output_item_done_event: Final = self.create_reasoning_output_item_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + self._pending_response_events.extend([text_done_event, part_done_event, reasoning_output_item_done_event]) + self._reasoning_done_emitted = True + self._reasoning_active = False + async def __anext__( self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: @@ -1078,7 +1138,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: - chunk = cast(ModelResponseStream, chunk) + chunk = self._reasoning_part_first(cast(ModelResponseStream, chunk)) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1095,55 +1155,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(chunk) ) - if self._reasoning_active and not self._reasoning_done_emitted: - # Incrementally accumulate reasoning content instead of - # calling stream_chunk_builder on every chunk (O(n²)) - delta = chunk.choices[0].delta if chunk.choices else None - if delta and hasattr(delta, "reasoning_content") and delta.reasoning_content: - self._accumulated_reasoning_content_parts.append(delta.reasoning_content) - if self._is_reasoning_end(chunk): - reasoning_content = "".join(self._accumulated_reasoning_content_parts) - - # Ensure we have a valid reasoning_item_id - self._cached_reasoning_item_id = ( - self._reasoning_item_id - or self._cached_reasoning_item_id - or mint_reasoning_item_id() - ) - reasoning_item_id = self._cached_reasoning_item_id - - # Create text.done event first with its own sequence number - self._sequence_number += 1 - text_done_event = self.create_reasoning_summary_text_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - - # Create part.done event second with its own sequence number - self._sequence_number += 1 - part_done_event = self.create_reasoning_summary_part_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - - self._sequence_number += 1 - reasoning_output_item_done_event = self.create_reasoning_output_item_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number, - ) - self._pending_response_events.extend( - [ - text_done_event, - part_done_event, - reasoning_output_item_done_event, - ] - ) - self._reasoning_done_emitted = True - self._reasoning_active = False - + self._queue_reasoning_done_events_on_reasoning_end(chunk) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) @@ -1182,11 +1194,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: buffered_chunk = self._take_buffered_chunk() if buffered_chunk is not None: - chunk = buffered_chunk + chunk = self._reasoning_part_first(buffered_chunk) elif self._upstream_exhausted: raise StopIteration else: - chunk = self.litellm_custom_stream_wrapper.__next__() + chunk = self._reasoning_part_first(self.litellm_custom_stream_wrapper.__next__()) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1206,6 +1218,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) + self._queue_reasoning_done_events_on_reasoning_end(chunk) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: self._pending_response_events.append(response_api_chunk) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 587cdfcee25..35d4b8bf1ba 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1008,6 +1008,7 @@ class LiteLLMCompletionResponsesConfig: instructions: str | None, session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), + previous_response_client_ids: tuple[str, ...] = (), ) -> dict: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1022,6 +1023,7 @@ class LiteLLMCompletionResponsesConfig: lookup: Final = session_lookup or partial( ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id, item_cache_control=keeps_item_marks, + client_response_ids=previous_response_client_ids, ) chat_completion_session: Final = ( await lookup(previous_response_id) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f566c8f4e02..b7385c3b3d5 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -16,6 +16,7 @@ from typing_extensions import ( ) import litellm +from litellm._internal_context import internal_sub_call from litellm._logging import verbose_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, @@ -29,6 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, update_responses_tools_with_model_file_ids, @@ -41,6 +43,11 @@ from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.round_billing import ( + billed_for_every_round, + first_round_billed_on_failure, + summed_cost_breakdown, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PromptObject, @@ -64,7 +71,7 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import all_litellm_params +from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, all_litellm_params from litellm.utils import ( ProviderConfigManager, client, @@ -161,6 +168,67 @@ def mock_responses_api_response( ) +_STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _client_trace_id(request_kwargs: Mapping[str, object]) -> str | None: + trace_id: Final = request_kwargs.get("litellm_trace_id") + return trace_id if isinstance(trace_id, str) else None + + +def _request_call_id(request_kwargs: Mapping[str, object]) -> str | None: + call_id: Final = request_kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def _logging_params(logging_obj: LiteLLMLoggingObj) -> Mapping[str, object]: + return _STR_MAPPING.validate_python(logging_obj.litellm_params) + + +def _request_logging_params( + logging_obj: LiteLLMLoggingObj, request_kwargs: Mapping[str, object] +) -> Mapping[str, object]: + logged_metadata: Final = ( + request_kwargs["litellm_metadata"] if "litellm_metadata" in request_kwargs else request_kwargs.get("metadata") + ) + return {**_logging_params(logging_obj), **request_kwargs, "metadata": logged_metadata} + + +def _request_trace_id(request_kwargs: Mapping[str, object]) -> str | None: + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLoggingObj): + return StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, litellm_params=_logging_params(logging_obj) + ) + return _client_trace_id(request_kwargs) + + +def _langfuse_request_trace_id(request_kwargs: Mapping[str, object], standard_trace_id: str | None) -> str | None: + from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata + + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLoggingObj): + return None + litellm_params: Final = _request_logging_params(logging_obj, request_kwargs) + trace_id: Final = langfuse_trace_id( + metadata=langfuse_trace_metadata(litellm_params.get("metadata"), litellm_params.get("proxy_server_request")), + standard_trace_id=standard_trace_id, + litellm_trace_id=logging_obj.litellm_trace_id, + litellm_call_id=logging_obj.litellm_call_id, + proxy_server_request=litellm_params.get("proxy_server_request"), + ) + return trace_id if isinstance(trace_id, str) else None + + +def _listing_trace_id(request_kwargs: Mapping[str, object], *, bridged_first_turn: bool) -> str | None: + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + if bridged_first_turn and isinstance(logging_obj, LiteLLMLoggingObj): + return StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, litellm_params=_request_logging_params(logging_obj, request_kwargs) + ) + return _request_trace_id(request_kwargs) + + async def aresponses_api_with_mcp( input: str | ResponseInputParam, model: str, @@ -223,6 +291,22 @@ async def aresponses_api_with_mcp( else (None, None, None, None) ) + request_kwargs: Final = _STR_MAPPING.validate_python(kwargs) + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + joins_request_trace: Final = isinstance(logging_obj, LiteLLMLoggingObj) + bridged_first_turn: Final = ( + joins_request_trace + and not stream + and previous_response_id is None + and not _provider_stores_the_first_round(model, custom_llm_provider, kwargs) + ) + + listing_trace_id: Final = ( + _client_trace_id(request_kwargs) + if stream + else _listing_trace_id(request_kwargs, bridged_first_turn=bridged_first_turn) + ) + # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( original_mcp_tools, @@ -230,7 +314,8 @@ async def aresponses_api_with_mcp( ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_trace_id=listing_trace_id, + request_trace_id=None if stream else _langfuse_request_trace_id(request_kwargs, listing_trace_id), mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), @@ -312,13 +397,17 @@ async def aresponses_api_with_mcp( ######################################################### # Make initial response API call ######################################################### - response: Final = await aresponses( - input=input, - model=model, - tools=all_tools, - previous_response_id=previous_response_id, - **initial_call_params, - ) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.carry_failed_attempt_usage() + with internal_sub_call(): + response: Final = await aresponses( + input=input, + model=model, + tools=all_tools, + previous_response_id=previous_response_id, + **initial_call_params, + ) + first_round_breakdown: Final = logging_obj.cost_breakdown if isinstance(logging_obj, LiteLLMLoggingObj) else None verbose_logger.debug("Initial response %s", response) @@ -327,107 +416,116 @@ async def aresponses_api_with_mcp( # If auto-execute tools is True, then we need to execute the tool calls ######################################################### if should_auto_execute and isinstance(response, ResponsesAPIResponse): - tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response) + with first_round_billed_on_failure(logging_obj, response): + tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response) - if tool_calls: - user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") + if tool_calls: + user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") - # Extract MCP auth headers from the request to pass to MCP server - ( - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers_from_request, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=secret_fields, - tools=tools, - ) - - tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map=tool_server_map, - served_tools=original_mcp_tools, - tool_calls=tool_calls, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers_from_request, - litellm_call_id=kwargs.get("litellm_call_id"), - litellm_trace_id=kwargs.get("litellm_trace_id"), - request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), - guardrail_context=MCPRequestContext.resolve_guardrail_context( - MappingProxyType({**kwargs, "metadata": metadata, "model": model}) - ), - ) - - if tool_results: - persistence_disabled: Final = LiteLLM_Proxy_MCP_Handler.is_persistence_disabled(call_params) - - follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=response, - tool_results=tool_results, - original_input=input, - preserve_reasoning=persistence_disabled, + # Extract MCP auth headers from the request to pass to MCP server + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers_from_request, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=secret_fields, + tools=tools, ) - # Prepare parameters for follow-up call (restores original stream setting) - follow_up_call_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( - call_params=call_params, original_stream_setting=stream or False + tool_trace_id: Final = _request_trace_id(request_kwargs) + tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map=tool_server_map, + served_tools=original_mcp_tools, + tool_calls=tool_calls, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers_from_request, + litellm_call_id=_request_call_id(request_kwargs), + litellm_trace_id=tool_trace_id, + request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), + guardrail_context=MCPRequestContext.resolve_guardrail_context( + MappingProxyType({**kwargs, "metadata": metadata, "model": model}) + ), + request_trace_id=_langfuse_request_trace_id(request_kwargs, tool_trace_id), ) - # Create tool execution events for streaming if needed - tool_execution_events = [] - if stream: - tool_execution_events = LiteLLM_Proxy_MCP_Handler.create_tool_execution_events( - tool_calls=tool_calls, tool_results=tool_results - ) + if tool_results: + chains_on_first_round: Final = not LiteLLM_Proxy_MCP_Handler.is_persistence_disabled( + call_params + ) and _provider_stores_the_first_round(model, custom_llm_provider, kwargs) - final_response = await LiteLLM_Proxy_MCP_Handler.make_follow_up_call( - follow_up_input=follow_up_input, - model=model, - all_tools=all_tools, - response_id=previous_response_id if persistence_disabled else response.id, - **follow_up_call_params, - ) - - # If streaming and we have tool execution events, wrap the response - if ( - stream - and tool_execution_events - and (hasattr(final_response, "__aiter__") or hasattr(final_response, "__iter__")) - ): - from litellm.responses.mcp.mcp_streaming_iterator import ( - MCPEnhancedStreamingIterator, - ) - - final_response = MCPEnhancedStreamingIterator( - tool_server_map=tool_server_map, - served_tools=original_mcp_tools, - base_iterator=final_response, - mcp_events=tool_execution_events, - user_api_key_auth=user_api_key_auth, - ) - - # Add custom output elements to the final response (for non-streaming) - elif isinstance(final_response, ResponsesAPIResponse): - # Fetch MCP tools again for output elements (without OpenAI transformation) - ( - mcp_tools_for_output, - _, - ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( - user_api_key_auth=user_api_key_auth, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - request_tags=LiteLLM_Proxy_MCP_Handler.get_parent_request_tags(kwargs), - raw_headers=discovery_raw_headers, - ) - final_response = LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( - response=final_response, - mcp_tools_fetched=mcp_tools_for_output, + follow_up_input: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=response, tool_results=tool_results, + original_input=input, + preserve_reasoning=not chains_on_first_round, ) - return final_response + + # Prepare parameters for follow-up call (restores original stream setting) + follow_up_call_params: Final = LiteLLM_Proxy_MCP_Handler.prepare_follow_up_call_params( + call_params=call_params, original_stream_setting=stream or False + ) + + # Create tool execution events for streaming if needed + tool_execution_events = [] + if stream: + tool_execution_events = LiteLLM_Proxy_MCP_Handler.create_tool_execution_events( + tool_calls=tool_calls, tool_results=tool_results + ) + + with internal_sub_call(): + final_response = await LiteLLM_Proxy_MCP_Handler.make_follow_up_call( + follow_up_input=follow_up_input, + model=model, + all_tools=all_tools, + response_id=response.id if chains_on_first_round else previous_response_id, + **follow_up_call_params, + ) + + # If streaming and we have tool execution events, wrap the response + if ( + stream + and tool_execution_events + and (hasattr(final_response, "__aiter__") or hasattr(final_response, "__iter__")) + ): + from litellm.responses.mcp.mcp_streaming_iterator import ( + MCPEnhancedStreamingIterator, + ) + + final_response = MCPEnhancedStreamingIterator( + tool_server_map=tool_server_map, + served_tools=original_mcp_tools, + base_iterator=final_response, + mcp_events=tool_execution_events, + user_api_key_auth=user_api_key_auth, + ) + + # Add custom output elements to the final response (for non-streaming) + elif isinstance(final_response, ResponsesAPIResponse): + final_round_breakdown: Final = ( + logging_obj.cost_breakdown if isinstance(logging_obj, LiteLLMLoggingObj) else None + ) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.cost_breakdown = summed_cost_breakdown( + first=response, + first_breakdown=first_round_breakdown, + final=final_response, + final_breakdown=final_round_breakdown, + ) + return billed_for_every_round( + first=response, + first_breakdown=first_round_breakdown, + final=LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + response=final_response, + mcp_tools_fetched=original_mcp_tools, + tool_results=tool_results, + ), + final_breakdown=final_round_breakdown, + ) + return final_response return response @@ -442,6 +540,18 @@ def _bridges_to_chat_completions( _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED: Final = frozenset({"client_metadata"}) +def _previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]: + """The forms of the previous response's id the client held, which its spend log row can be stored under.""" + if not previous_response_id: + return () + addressed_id: Final = kwargs.get(ADDRESSED_RESPONSE_ID_FIELD) + return ( + (previous_response_id, addressed_id) + if isinstance(addressed_id, str) and addressed_id + else (previous_response_id,) + ) + + def _bridge_kwargs( kwargs: Mapping[str, object], responses_api_provider_config: BaseResponsesAPIConfig | None, @@ -569,6 +679,18 @@ def _will_bridge_to_chat_completions( ) +def _provider_stores_the_first_round(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool: + """A bridged first MCP round runs as an internal call, so no spend log holds it for the follow-up to replay.""" + api_base: Final = _api_base_kwarg(kwargs) + return not _will_bridge_to_chat_completions( + model, + custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1], + bool(kwargs.get("use_chat_completions_api")), + kwargs.get("model_info"), + api_base, + ) + + def _is_a_client_input_item( item: object, ) -> TypeIs[Mapping[str, object]]: # guard-ok: bounds a client input item before its mark is read @@ -1483,6 +1605,7 @@ def responses( if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api): bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params) return litellm_completion_transformation_handler.response_api_handler( + _previous_response_client_ids(previous_response_id, kwargs), model=model, input=input, responses_api_request=response_api_optional_params, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index b711d051055..b4dbf4c3d4e 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -258,6 +258,7 @@ class LiteLLM_Proxy_MCP_Handler: request_tags: list[str] | None = None, raw_headers: dict[str, str] | None = None, granted_toolsets: "GrantedToolsetIds | None" = None, + request_trace_id: str | None = None, ) -> tuple[list[MCPTool], list[str]]: """ Get available tools from the MCP server manager. @@ -343,6 +344,7 @@ class LiteLLM_Proxy_MCP_Handler: log_list_tools_to_spendlogs=True, list_tools_log_source="responses", litellm_trace_id=litellm_trace_id, + request_trace_id=request_trace_id, request_tags=request_tags, raw_headers=raw_headers, ) @@ -474,6 +476,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, request_tags: list[str] | None = None, raw_headers: dict[str, str] | None = None, + request_trace_id: str | None = None, ) -> tuple[list[MCPTool], dict[str, str]]: """ Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. @@ -505,6 +508,7 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers=mcp_server_auth_headers, request_tags=request_tags, raw_headers=raw_headers, + request_trace_id=request_trace_id, ) # Step 2: Filter tools based on allowed_tools parameter @@ -729,6 +733,7 @@ class LiteLLM_Proxy_MCP_Handler: request_tags: list[str] | None = None, guardrail_context: Mapping[str, object] | None = None, served_tools: Sequence[MCPTool] | None = None, + request_trace_id: str | None = None, ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -787,7 +792,11 @@ class LiteLLM_Proxy_MCP_Handler: }, } ] - tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + tool_logging_call_id = ( + f"{litellm_call_id}_{tool_call_id or uuid.uuid4()}" + if request_trace_id and litellm_call_id + else litellm_call_id or str(uuid.uuid4()) + ) logging_metadata: dict[str, object] = { "tool_call_id": tool_call_id, "tool_name": sanitized_tool_name, @@ -813,6 +822,8 @@ class LiteLLM_Proxy_MCP_Handler: } if litellm_trace_id: logging_request_data["litellm_trace_id"] = litellm_trace_id + if request_trace_id: + logging_metadata["existing_trace_id"] = request_trace_id if request_tags: logging_metadata["tags"] = request_tags if typed_user_api_key_auth is not None: diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py new file mode 100644 index 00000000000..f20d01785a5 --- /dev/null +++ b/litellm/responses/mcp/round_billing.py @@ -0,0 +1,134 @@ +from asyncio import CancelledError +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.hidden_params import served_from_cache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import CostBreakdown + +_PER_REQUEST_KEYS: Final = frozenset({"discount_percent", "margin_percent", "margin_fixed_amount"}) +_TOTALS_WITH_FIXED_MARGIN: Final = ("total_cost", "margin_total_amount") +_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown) +_STR_KEYED: Final = TypeAdapter(Mapping[str, object]) + + +def _summed(first: object, final: object) -> object: + if first is None: + return final + if final is None: + return first + if isinstance(first, bool) or isinstance(final, bool): + return final + if isinstance(first, (int, float)) and isinstance(final, (int, float)): + return first + final + if isinstance(first, Mapping) and isinstance(final, Mapping): + return _summed_mapping(_STR_KEYED.validate_python(first), _STR_KEYED.validate_python(final)) + return final + + +def _summed_entry(key: str, first: Mapping[str, object], final: Mapping[str, object]) -> object: + if key in _PER_REQUEST_KEYS: + return final[key] if key in final else first[key] + return _summed(first.get(key), final.get(key)) + + +def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> Mapping[str, object]: + keys: Final = (*first, *(key for key in final if key not in first)) + return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) + + +def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: + if first is None or final is None: + return final if first is None else first + return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + + +def _repeated_fixed_margin( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> float: + if first_breakdown is None or final_breakdown is None or first_breakdown is final_breakdown: + return 0.0 + if served_from_cache(first) or served_from_cache(final): + return 0.0 + return min(first_breakdown.get("margin_fixed_amount", 0.0), final_breakdown.get("margin_fixed_amount", 0.0)) + + +def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) -> Mapping[str, object]: + charged_once: Final = MappingProxyType( + {key: total - repeated for key in _TOTALS_WITH_FIXED_MARGIN if isinstance(total := summed.get(key), float)} + ) + return MappingProxyType({**summed, **charged_once}) + + +def summed_cost_breakdown( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> CostBreakdown | None: + """The breakdown of the charged rounds, or None when a charged round was priced without its own breakdown, since a + partial one would disagree with the billed cost. The fixed margin is charged once per request, not once per round.""" + if served_from_cache(final): + return None if served_from_cache(first) else first_breakdown + if first_breakdown is final_breakdown: + return None + if served_from_cache(first): + return final_breakdown + if first_breakdown is None or final_breakdown is None: + return None + return _COST_BREAKDOWN.validate_python( + _less_repeated_fixed_margin( + _summed_mapping(first_breakdown, final_breakdown), + _repeated_fixed_margin(first, first_breakdown, final, final_breakdown), + ) + ) + + +def billed_for_every_round( + first: ResponsesAPIResponse, + first_breakdown: CostBreakdown | None, + final: ResponsesAPIResponse, + final_breakdown: CostBreakdown | None, +) -> ResponsesAPIResponse: + """The final round's response carrying the usage and cost of both rounds, so the one logged row bills the request. + + Costs are summed per round rather than repriced from the summed usage, because a round's price can depend on its + own size (tiered rates) or on charges the usage does not show (built-in tools, provider-reported costs). The fixed + margin both rounds added is charged once, since it is a per-request fee. + """ + first_cost: Final = first.hidden_params.get("response_cost") + final_cost: Final = final.hidden_params.get("response_cost") + total_cost: Final = ( + first_cost + final_cost - _repeated_fixed_margin(first, first_breakdown, final, final_breakdown) + if isinstance(first_cost, (int, float)) and isinstance(final_cost, (int, float)) + else None + ) + billed: Final = final.model_copy(update=MappingProxyType({"usage": merged_round_usage(first.usage, final.usage)})) + billed.hidden_params = { + **final.hidden_params, + "response_cost": total_cost, + } + return billed + + +@contextmanager +def first_round_billed_on_failure(logging_obj: object, first: ResponsesAPIResponse) -> Generator[None]: + try: + yield + except (Exception, CancelledError): + first_cost: Final = first.hidden_params.get("response_cost") + if isinstance(logging_obj, LiteLLMLoggingObj) and first.usage is not None and not served_from_cache(first): + logging_obj.record_partial_usage_for_failure( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(first.usage), + float(first_cost) if isinstance(first_cost, (int, float)) else 0.0, + ) + raise diff --git a/litellm/router.py b/litellm/router.py index 3c969610e83..99adc6c5fa6 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 ( @@ -500,7 +501,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 @@ -537,6 +540,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]) @@ -573,6 +577,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 {} @@ -4507,7 +4526,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). @@ -4518,11 +4539,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"] @@ -4543,7 +4565,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"] @@ -6035,6 +6057,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, ) @@ -6044,6 +6067,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 @@ -6130,11 +6161,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 @@ -6150,6 +6181,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 @@ -6169,6 +6201,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, @@ -10964,6 +10997,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 @@ -13727,6 +13761,15 @@ class Router: """ return [m for m in self.model_list if m["litellm_params"]["model"] == model] + def get_routable_deployments(self, model: str) -> tuple[DeploymentTypedDict, ...]: + try: + _, candidates = self._common_checks_available_deployment(model=model, request_kwargs={}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped + except (litellm.BadRequestError, ValueError): + return () + if isinstance(candidates, dict): + return (cast(DeploymentTypedDict, candidates),) # cast-ok: the router returns one deployment as a dict + return tuple(cast(list[DeploymentTypedDict], candidates)) # cast-ok: the router returns deployments as a list + def _try_early_resolve_deployments_for_model_not_in_names( self, model: str, diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b20f71cd51f..1622a30b95f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -405,7 +405,8 @@ class RouterBudgetLimiting(CustomLogger): Runs once the budget start time exists in Redis Cache (on the 2nd and subsequent requests to the same provider) - Increments the spend in memory cache (so spend instantly updated in memory) - - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) + - Queues the increment for the batched Redis pipeline when Redis is configured, since only a multi instance + deployment has anything to sync """ increment_op: Final = RedisPipelineIncrementOperation( key=spend_key, @@ -418,6 +419,8 @@ class RouterBudgetLimiting(CustomLogger): value=response_cost, ttl=ttl, ) + if self.dual_cache.redis_cache is None: + return self.redis_increment_operation_queue.append(increment_op) def _get_redis_increment_queue_lock(self) -> asyncio.Lock: 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/llms/anthropic.py b/litellm/types/llms/anthropic.py index c3d257117ed..6aeead7679f 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -536,7 +536,9 @@ class ContentBlockStartText(TypedDict): content_block: TextBlock -ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinkingBlock +ContentBlockContentBlockDict = ( + ToolUseBlock | TextBlock | ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock +) ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 72e564b6919..f0be66f055d 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -248,6 +248,15 @@ class MCPCredentials(TypedDict, total=False): Optional RFC 8707 resource indicator sent on ID-JAG leg 1 """ + id_jag_subject_source: ReadOnly[str | None] + """ + Which token ID-JAG leg 1 presents as its RFC 8693 subject_token. ``login_assertion`` (the default) is the + identity assertion captured at SSO login, an OIDC id_token or a SAML assertion. ``refresh_token`` exchanges + that assertion at the IdP token endpoint once, at login, for a refresh token and presents the refresh token + on every leg 1 instead, which is how a SAML requesting app works at IdPs whose ID-JAG exchange does not take + a saml2 subject directly. + """ + id_jag_resource_client_id: ReadOnly[str | None] """OAuth client identifier used at the ID-JAG resource authorization server.""" diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index f7c4d329781..219f23200d2 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -182,6 +182,7 @@ class MCPServer(LiteLLMBaseModel): # posts the ID-JAG assertion to id_jag_resource_token_endpoint. id_jag_resource_token_endpoint: str | None = None id_jag_resource: str | None = None + id_jag_subject_source: str | None = None id_jag_resource_client_id: str | None = Field( default=None, description="OAuth client identifier used at the ID-JAG resource authorization server.", 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 7df8cc7452f..c9925580190 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2,6 +2,7 @@ import json import re import time from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass from enum import Enum from types import MappingProxyType from typing import ( @@ -3899,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 @@ -3926,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 @@ -4268,6 +4269,7 @@ class LlmProviders(str, Enum): DEEPINFRA = "deepinfra" PERPLEXITY = "perplexity" TYPESAFE = "typesafe" + NACE = "nace" STRANDS_DECIDER = "strands_decider" LAYA = "laya" BESPOKE = "bespoke" @@ -4477,6 +4479,12 @@ class LiteLLMLoggingBaseClass: TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] +@dataclass(frozen=True, slots=True) +class CountContextManagement: + spec: JsonValue + drop_invalid_edits: bool + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/litellm/utils.py b/litellm/utils.py index d2ca9faef41..04c1e628288 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, @@ -1374,15 +1376,16 @@ def _dispatch_success_logging( is_completion_with_fallbacks: bool, is_litellm_internal_call: bool, ) -> None: - if not is_litellm_internal_call: - _schedule_async_success_logging( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + if is_litellm_internal_call: + return + _schedule_async_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, @@ -1494,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. @@ -1505,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: @@ -3417,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. @@ -3445,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 = {} @@ -3530,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 @@ -9116,6 +9103,8 @@ class ProviderConfigManager: return litellm.PerplexityDecisionsConfig() if provider == LlmProviders.TYPESAFE: return litellm.TypeSafeDecisionsConfig() + if provider == LlmProviders.NACE: + return litellm.NaceDecisionsConfig() if provider == LlmProviders.OPENROUTER: return litellm.OpenRouterDecisionsConfig() if provider == LlmProviders.FIREWORKS_AI: @@ -9251,7 +9240,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() @@ -9975,6 +9964,7 @@ class ProviderConfigManager: def get_provider_realtime_config( model: str, provider: LlmProviders, + drop_params: bool = False, ) -> BaseRealtimeConfig | None: if LlmProviders.GEMINI == provider: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig @@ -9983,7 +9973,7 @@ class ProviderConfigManager: if LlmProviders.META == provider: from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig - return MetaRealtimeConfig() + return MetaRealtimeConfig(drop_params=drop_params) return None @staticmethod 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..1851e818ec7 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, @@ -70917,6 +71184,38 @@ "notes": "Self-hosted decision model; infrastructure costs are paid separately" } }, + "nace/drex-v1.0": { + "input_cost_per_token": 4e-08, + "litellm_provider": "nace", + "max_input_tokens": 32768, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v1.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-v2.0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, + "nace/drex-latest": { + "input_cost_per_token": 6e-08, + "litellm_provider": "nace", + "max_input_tokens": 131072, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://console.nace.ai/docs/drex/reference/models" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -75218,5 +75517,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/provider_endpoints_support.json b/provider_endpoints_support.json index 0325fe1f796..ad38b064903 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2625,6 +2625,13 @@ "image_variations": true } }, + "nace": { + "display_name": "Nace (`nace`)", + "url": "https://console.nace.ai/docs/drex", + "endpoints": { + "systemone": true + } + }, "typesafe": { "display_name": "TypeSafe (`typesafe`)", "url": "https://docs.typesafe.ai/models", 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 6762e4cd35c..e669107900d 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/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 3ad80af10f0..f50096add19 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -78,6 +78,7 @@ - {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"} - {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"} - {id: mgmt.fallback_management.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "fallback_management_endpoints.py", rationale: "Fallback config (smoke)"} +- {id: mgmt.fallback_management.update.config_owned_refused, module: mgmt, tier: P1, surface: api, assertions: [config_owned_refused], source: "fallback_management_endpoints.py", rationale: "A /fallback write to a fallbacks key the config file sets is refused and the config rule keeps routing"} - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} diff --git a/tests/e2e/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/management/test_config_misc_endpoints_e2e.py b/tests/e2e/management/test_config_misc_endpoints_e2e.py index b7ab311cc18..bd9192fa0f0 100644 --- a/tests/e2e/management/test_config_misc_endpoints_e2e.py +++ b/tests/e2e/management/test_config_misc_endpoints_e2e.py @@ -31,7 +31,7 @@ import pytest from pydantic import BaseModel from e2e_config import unique_marker -from e2e_http import NoBody, Success, unwrap, unwrap_status +from e2e_http import NoBody, Result, Success, UnknownApiError, unwrap, unwrap_status from e2e_metadata import Domain, Route, Subject, meta from lifecycle import ResourceManager from management_client import ManagementClient @@ -141,12 +141,15 @@ class ComplianceResponse(BaseModel): # ---- fallback management --------------------------------------------------- -class FallbackShape(BaseModel): +class FallbackTarget(BaseModel): model: str - fallback_models: list[str] fallback_type: str +class FallbackShape(FallbackTarget): + fallback_models: list[str] + + class FallbackCreateBody(FallbackShape): pass @@ -163,6 +166,28 @@ class FallbackGetResponse(FallbackShape): pass +class FallbackDeleteResponse(FallbackTarget): + message: str + + +class RouterFallbackValues(BaseModel): + fallbacks: list[dict[str, list[str]]] | None = None + + +class RouterFallbackOwnership(BaseModel): + current_values: RouterFallbackValues + source: dict[str, str] + + +class ConfigOwnedDetail(BaseModel): + keys: list[str] + section: str + + +class ConfigOwnedRejection(BaseModel): + detail: ConfigOwnedDetail + + # ---- jwt key mapping ------------------------------------------------------- @@ -365,20 +390,19 @@ class TestFallbackManagement: @pytest.mark.covers("mgmt.fallback_management.update.happy_path") @meta(Subject(domain=Domain.ROUTING)) def test_create_persists_and_is_read_back(self, client: ManagementClient, resources: ResourceManager) -> None: - primary = f"e2e-fallback-primary-{unique_marker()}" - secondary = f"e2e-fallback-secondary-{unique_marker()}" - params = LiteLLMParamsBody(model="openai/gpt-5.5", api_key="e2e-dummy-key") - primary_id = client.proxy.create_model(primary, params) - resources.defer(lambda: client.proxy.delete_model(primary_id)) - secondary_id = client.proxy.create_model(secondary, params) - resources.defer(lambda: client.proxy.delete_model(secondary_id)) - resources.defer(lambda: self._delete_fallback(client, primary)) + ownership = self._router_fallback_ownership(client) + assert ownership.source.get("context_window_fallbacks") != "config", ( + "the proxy config sets router_settings.context_window_fallbacks, so /fallback cannot write it and this " + "round trip needs a fallback type the database owns" + ) + primary, secondary = self._register_fallback_pair(client, resources) + resources.defer(lambda: self._delete_fallback(client, primary, "context_window")) created = unwrap( client.proxy.transport.post( "/fallback", headers=client.proxy.transport.master, - json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="general"), + json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="context_window"), response_type=FallbackResponse, ) ) @@ -388,13 +412,7 @@ class TestFallbackManagement: ) def read_back() -> FallbackGetResponse | None: - result = client.proxy.transport.get( - f"/fallback/{primary}", - headers=client.proxy.transport.master, - params=FallbackGetParams(fallback_type="general"), - response_type=FallbackGetResponse, - ) - match result: + match self._get_fallback(client, primary, "context_window"): case Success(data=data) if secondary in data.fallback_models: return data case _: @@ -405,14 +423,105 @@ class TestFallbackManagement: f"GET /fallback/{primary} reports fallbacks {got.fallback_models}, configured [{secondary!r}]" ) + deleted = unwrap(self._delete_fallback(client, primary, "context_window")) + assert deleted.model == primary and deleted.fallback_type == "context_window", ( + f"DELETE /fallback/{primary} answered for model={deleted.model!r} type={deleted.fallback_type!r}" + ) + + def gone() -> bool | None: + match self._get_fallback(client, primary, "context_window"): + case UnknownApiError(status_code=404): + return True + case _: + return None + + _ = _poll(client, gone, f"GET /fallback/{primary} still reports a fallback after DELETE /fallback/{primary}") + + @pytest.mark.covers("mgmt.fallback_management.update.config_owned_refused") + @meta(Subject(domain=Domain.ROUTING)) + def test_write_to_config_owned_fallbacks_is_refused( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + before = self._router_fallback_ownership(client) + config_rules = before.current_values.fallbacks + assert before.source.get("fallbacks") == "config" and config_rules, ( + "this test needs router_settings.fallbacks set in the proxy config, but /router/settings reports " + f"source={before.source.get('fallbacks')!r} value={config_rules}" + ) + primary, secondary = self._register_fallback_pair(client, resources) + resources.defer(lambda: self._delete_fallback(client, primary, "general")) + + result = client.proxy.transport.post( + "/fallback", + headers=client.proxy.transport.master, + json=FallbackCreateBody(model=primary, fallback_models=[secondary], fallback_type="general"), + response_type=FallbackResponse, + ) + match result: + case UnknownApiError(status_code=400, body=body): + rejection = ConfigOwnedRejection.model_validate_json(body).detail + case _: + pytest.fail(f"/fallback accepted a write to the config-owned router_settings.fallbacks: {result}") + assert rejection.keys == ["fallbacks"] and rejection.section == "router_settings", ( + f"/fallback refused the write naming keys={rejection.keys} section={rejection.section!r}, " + "expected router_settings.fallbacks" + ) + + after = self._router_fallback_ownership(client).current_values.fallbacks + assert after == config_rules, ( + f"the refused write changed router_settings.fallbacks from {config_rules} to {after}" + ) + match self._get_fallback(client, primary, "general"): + case UnknownApiError(status_code=404): + pass + case refused_rule: + pytest.fail(f"the refused rule for {primary} is served anyway: {refused_rule}") + config_model, config_targets = next(iter(config_rules[0].items())) + served = unwrap(self._get_fallback(client, config_model, "general")) + assert served.fallback_models == config_targets, ( + f"GET /fallback/{config_model} reports {served.fallback_models}, " + f"the config file routes it to {config_targets}" + ) + @staticmethod - def _delete_fallback(client: ManagementClient, model: str) -> None: - _ = client.proxy.transport.delete( + def _register_fallback_pair(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]: + primary = f"e2e-fallback-primary-{unique_marker()}" + secondary = f"e2e-fallback-secondary-{unique_marker()}" + params = LiteLLMParamsBody(model="openai/gpt-5.5", api_key="e2e-dummy-key") + primary_id = client.proxy.create_model(primary, params) + resources.defer(lambda: client.proxy.delete_model(primary_id)) + secondary_id = client.proxy.create_model(secondary, params) + resources.defer(lambda: client.proxy.delete_model(secondary_id)) + return primary, secondary + + @staticmethod + def _router_fallback_ownership(client: ManagementClient) -> RouterFallbackOwnership: + return unwrap( + client.proxy.transport.get( + "/router/settings", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=RouterFallbackOwnership, + ) + ) + + @staticmethod + def _get_fallback(client: ManagementClient, model: str, fallback_type: str) -> Result[FallbackGetResponse]: + return client.proxy.transport.get( + f"/fallback/{model}", + headers=client.proxy.transport.master, + params=FallbackGetParams(fallback_type=fallback_type), + response_type=FallbackGetResponse, + ) + + @staticmethod + def _delete_fallback(client: ManagementClient, model: str, fallback_type: str) -> Result[FallbackDeleteResponse]: + return client.proxy.transport.delete( f"/fallback/{model}", headers=client.proxy.transport.master, json=NoBody(), - params=FallbackGetParams(fallback_type="general"), - response_type=NoBody, + params=FallbackGetParams(fallback_type=fallback_type), + response_type=FallbackDeleteResponse, ) 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/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py index 41b61cdc987..89263977fb0 100644 --- a/tests/e2e/perf/cases/mixed_size_requests_2rps.py +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -69,6 +69,9 @@ CASE: Final = PerformanceTestCase( Limit("memory_growth_megabytes", max=600), Limit("cpu_seconds_per_request", max=1.26), Limit("postgres_calls_per_request", max=3.5), + Limit("postgres_max_statement_per_request", max=0.5), + Limit("postgres_max_background_statement_per_request", max=0.9), Limit("redis_calls_per_request", max=21.0), ), + background_cap_exempt_tables=frozenset({"LiteLLM_DailyGuardrailMetrics"}), ) diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index 00a9fdefdb5..73772465492 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -68,6 +68,8 @@ CASE: Final = PerformanceTestCase( Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), Limit("postgres_calls_per_request", max=1.8), + Limit("postgres_max_statement_per_request", max=0.5), + Limit("postgres_max_background_statement_per_request", max=0.9), Limit("redis_calls_per_request", max=18), Limit("cpu_seconds_per_request", max=0.29), ), diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index b4ad10b5a5a..54c6943a054 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -76,6 +76,11 @@ class StatementCalls: return "plumbing" return "awaited" if verb == "SELECT" and tables[0] in _AWAITED_TABLES else "background" + @property + def table(self) -> str: + tables: Final = _TABLE.findall(self.query) + return tables[0] if tables else "" + @property def summary(self) -> str: return re.sub(r"\s+", " ", self.query)[:110] diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index 98391a02418..f634df7de84 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -128,6 +128,21 @@ class Observation: def statements(self) -> tuple[StatementCalls, ...] | None: return statement_deltas(self.db_before, self.db_after) + @property + def busiest_awaited_statement(self) -> StatementCalls | None: + """The awaited statement with the most calls: a lookup added to every request shows up here at + 1.0 per request while cache refreshes stay fractional, which the awaited total cannot tell apart.""" + awaited: Final = tuple(item for item in self.statements or () if item.bucket == "awaited") + return max(awaited, key=_item_count, default=None) + + def busiest_background_statement(self, exempt_tables: frozenset[str]) -> StatementCalls | None: + """The background statement with the most calls, skipping the case's per-flush exemptions: a rollup + written once per request reads 1.0 here while batched writers stay below it.""" + background: Final = tuple( + item for item in self.statements or () if item.bucket == "background" and item.table not in exempt_tables + ) + return max(background, key=_item_count, default=None) + @property def retried_markers(self) -> frozenset[str]: counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) @@ -195,7 +210,9 @@ def _text_sha1(text: str, *, intervened: bool) -> str: return hashlib.sha1(original.encode()).hexdigest() -def observe(metric: Metric, percentile_rank: int | None, observation: Observation) -> float | None: +def observe( + metric: Metric, percentile_rank: int | None, observation: Observation, exempt_tables: frozenset[str] = frozenset() +) -> float | None: """The observed value of one metric, at a percentile when the metric is per request.""" per_request: dict[str, list[float]] = { "gateway_added_seconds": observation.gateway_added_seconds(), @@ -209,6 +226,16 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio return observation.count_per_request(calls_in("awaited", observation.statements)) if metric == "postgres_background_calls_per_request": return observation.count_per_request(calls_in("background", observation.statements)) + if metric == "postgres_max_statement_per_request": + if observation.statements is None: + return None + busiest: Final = observation.busiest_awaited_statement + return observation.count_per_request(0 if busiest is None else busiest.calls) + if metric == "postgres_max_background_statement_per_request": + if observation.statements is None: + return None + busiest_background: Final = observation.busiest_background_statement(exempt_tables) + return observation.count_per_request(0 if busiest_background is None else busiest_background.calls) if metric == "redis_calls_per_request": return observation.count_per_request(calls_in("awaited", observation.redis_commands)) if metric == "redis_background_calls_per_request": @@ -225,6 +252,7 @@ class Row: percentile: int | None observed: float | None limit: float | None + detail: str = "" @property def label(self) -> str: @@ -252,6 +280,8 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("memory_growth_megabytes", None), ("postgres_calls_per_request", None), ("postgres_background_calls_per_request", None), + ("postgres_max_statement_per_request", None), + ("postgres_max_background_statement_per_request", None), ("redis_calls_per_request", None), ("redis_background_calls_per_request", None), ("cpu_seconds_per_request", None), @@ -273,12 +303,13 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ Row( metric=metric, percentile=rank, - observed=observe(metric, rank, observation), + observed=observe(metric, rank, observation, case.background_cap_exempt_tables), limit=limit.max if limit else None, + detail=_detail(metric, observation, case.background_cap_exempt_tables), ) ) exceeded: Final = [ - f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}" + f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}{f' ({row.detail})' if row.detail else ''}" if row.over and row.observed is not None and row.limit is not None else f"{row.label}: limit set but nothing observed" for row in rows @@ -287,6 +318,16 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ return rows, exceeded +def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[str]) -> str: + if metric == "postgres_max_statement_per_request": + busiest: Final = observation.busiest_awaited_statement + return "" if busiest is None else busiest.summary + if metric == "postgres_max_background_statement_per_request": + busiest_background: Final = observation.busiest_background_statement(exempt_tables) + return "" if busiest_background is None else busiest_background.summary + return "" + + def _row_properties(row: Row) -> tuple[tuple[str, float], ...]: observed: Final = ((row.property_name, row.observed),) if row.observed is not None else () limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else () diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 76e6efa3203..37ad640db62 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -23,6 +23,8 @@ Metric = Literal[ "memory_growth_megabytes", "postgres_calls_per_request", "postgres_background_calls_per_request", + "postgres_max_statement_per_request", + "postgres_max_background_statement_per_request", "cpu_seconds_per_request", "redis_calls_per_request", "redis_background_calls_per_request", @@ -148,6 +150,9 @@ class PerformanceTestCase: traffic: Traffic limits: tuple[Limit, ...] mock_latency: MockLatency = field(default_factory=lambda: MockLatency()) + background_cap_exempt_tables: frozenset[str] = frozenset() + """Tables whose background write is one upsert per writer flush and, at this case's request rate, lands + above the `postgres_max_background_statement_per_request` cap without being per request.""" def team(self, name: str) -> Team: return next(team for team in self.teams if team.name == name) 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/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts index 50850a5cd58..d0b1bb9ecfb 100644 --- a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -18,8 +18,11 @@ const test = base.extend<{ scopedKey: ScopedKey }>({ team_id: E2E_TEAM_CRUD_ID, models: [CHAT_MODEL_A], }); - await use({ alias, token: created.token }); - await deleteVirtualKey(page.request, created.token); + try { + await use({ alias, token: created.token }); + } finally { + await deleteVirtualKey(page.request, created.token); + } }, }); @@ -36,15 +39,44 @@ test.describe("Proxy Admin - Key budget window", () => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - await openKeyDetail(page, alias); - - await page.getByRole("tab", { name: "Settings" }).click(); - await page.getByRole("button", { name: "Edit Settings" }).click(); - - await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + // The initial, unfiltered list must already contain this newly created key. + // useKeys keeps that list visible while the filtered search is pending. + await expect(page.getByRole("row").filter({ hasText: alias })).toBeVisible({ timeout: 15_000 }); + // Deliver the real search result after editing starts: replacing the list's + // key object must not reset the open form and erase its unsaved budget. + const searchReceived = Promise.withResolvers(); + const releaseSearch = Promise.withResolvers(); + const isAliasSearch = (url: string) => { + const parsed = new URL(url); + return parsed.pathname.endsWith("/key/list") && parsed.searchParams.get("search") === alias; + }; + await page.route("**/key/list?**", async (route) => { + if (!isAliasSearch(route.request().url())) return route.continue(); + const response = await route.fetch(); + searchReceived.resolve(); + await releaseSearch.promise; + await route.fulfill({ response }); + }); + try { + await openKeyDetail(page, alias); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await searchReceived.promise; + const refreshed = page.waitForResponse((response) => isAliasSearch(response.url())); + releaseSearch.resolve(); + await refreshed; + } finally { + releaseSearch.resolve(); + await page.unrouteAll({ behavior: "wait" }); + } await page.getByLabel("Reset Budget", { exact: true }).click(); await page.getByRole("option", { name: "monthly", exact: true }).click(); - await page.getByRole("button", { name: "Save Changes" }).click(); + const saved = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(saved.max_budget, "a delayed search response must not erase the budget draft").toBe("12.5"); + expect(saved.budget_duration).toBe("30d"); await expect .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { 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/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py index f38e7927fb7..ad54c61aed0 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -69,5 +69,101 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_gateway_added_seconds_p99": 0.75, "perf_postgres_calls_per_request": 10.0, "perf_postgres_background_calls_per_request": 0.0, + "perf_postgres_max_statement_per_request": 10.0, + "perf_postgres_max_background_statement_per_request": 0.0, "perf_redis_calls_per_request_limit": 35.0, } + + +def _statement(queryid: int, query: str, calls: int) -> StatementCalls: + return StatementCalls(queryid=queryid, query=query, calls=calls) + + +def _observation_with(statements: tuple[StatementCalls, ...]) -> Observation: + before: Final = tuple(replace(item, calls=0) for item in statements) + return Observation( + records=tuple(served(seq, 1.0) for seq in range(1000)), + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=600.0, + db_before=DbCounters(commits=0, statements=before), + db_after=DbCounters(commits=0, statements=statements), + ) + + +MAX_STATEMENT_CASE: Final = replace(CASE, limits=(Limit("postgres_max_statement_per_request", max=0.5),)) +BACKGROUND_CAP_CASE: Final = replace( + CASE, + limits=(Limit("postgres_max_background_statement_per_request", max=0.9),), + background_cap_exempt_tables=frozenset({"LiteLLM_DailyGuardrailMetrics"}), +) +GUARDRAIL_UPSERT: Final = 'INSERT INTO "LiteLLM_DailyGuardrailMetrics" ("guardrail_id","date") VALUES ($1,$2)' +MODEL_USAGE_UPSERT: Final = 'INSERT INTO "LiteLLM_DailyModelUsage" ("date","model") VALUES ($1,$2)' +TEAM_REFRESH: Final = 'SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id = $1' +USER_LOOKUP: Final = 'SELECT * FROM "LiteLLM_UserTable" WHERE user_id = $1' + + +def test_cache_refreshes_whose_sum_exceeds_the_cap_pass_the_per_statement_limit() -> None: + statements: Final = ( + _statement(1, TEAM_REFRESH, calls=357), + _statement(2, USER_LOOKUP, calls=245), + _statement(3, 'INSERT INTO "LiteLLM_DailyModelUsage" VALUES ($1)', calls=1000), + ) + + rows, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements)) + + assert exceeded == [] + assert dict(junit_properties(rows))["perf_postgres_max_statement_per_request"] == 0.357 + assert dict(junit_properties(rows))["perf_postgres_calls_per_request"] == 0.602 + + +def test_a_lookup_on_every_request_fails_the_per_statement_limit_and_names_the_statement() -> None: + statements: Final = (_statement(1, TEAM_REFRESH, calls=357), _statement(2, USER_LOOKUP, calls=1000)) + + _, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements)) + + assert exceeded == [f"postgres_max_statement_per_request: observed 1.0000 > limit 0.5000 ({USER_LOOKUP})"], exceeded + + +def test_per_statement_limit_without_db_counters_reports_nothing_observed() -> None: + observation: Final = replace(_observation_with(()), db_before=None, db_after=None) + + _, exceeded = evaluate(MAX_STATEMENT_CASE, observation) + + assert exceeded == ["postgres_max_statement_per_request: limit set but nothing observed"] + + +def test_per_flush_rollups_pass_the_background_cap_and_the_exempt_table_is_skipped() -> None: + statements: Final = ( + _statement(1, GUARDRAIL_UPSERT, calls=1150), + _statement(2, MODEL_USAGE_UPSERT, calls=810), + _statement(3, USER_LOOKUP, calls=1000), + ) + + rows, exceeded = evaluate(BACKGROUND_CAP_CASE, _observation_with(statements)) + + assert exceeded == [] + assert dict(junit_properties(rows))["perf_postgres_max_background_statement_per_request"] == 0.81 + + +def test_a_rollup_written_on_every_request_fails_the_background_cap_and_names_the_statement() -> None: + statements: Final = (_statement(1, GUARDRAIL_UPSERT, calls=1050), _statement(2, MODEL_USAGE_UPSERT, calls=1000)) + + _, exceeded = evaluate(BACKGROUND_CAP_CASE, _observation_with(statements)) + + assert exceeded == [ + f"postgres_max_background_statement_per_request: observed 1.0000 > limit 0.9000 ({MODEL_USAGE_UPSERT})" + ], exceeded + + +def test_background_cap_without_exemptions_counts_the_guardrail_upsert() -> None: + case: Final = replace(BACKGROUND_CAP_CASE, background_cap_exempt_tables=frozenset()) + statements: Final = (_statement(1, GUARDRAIL_UPSERT, calls=1150), _statement(2, MODEL_USAGE_UPSERT, calls=810)) + + _, exceeded = evaluate(case, _observation_with(statements)) + + assert exceeded == [ + f"postgres_max_background_statement_per_request: observed 1.1500 > limit 0.9000 ({GUARDRAIL_UPSERT})" + ], exceeded 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/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py index d30bf1b92f5..ff618a00136 100644 --- a/tests/integration/_support/bedrock_realtime_peer.py +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -14,6 +14,12 @@ text message (the injected trigger-audio block is USER/AUDIO, so it never trips ``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` ``toolUseId`` names the call being answered. + +A ``Refusal`` peer withholds each stream's initial response until that stream's opening arrived +(its SYSTEM prompt block or its first USER text block closed), then answers the HTTP error Bedrock +sends for an unknown model id. A ``StreamFailure`` peer accepts the stream, plays its turns, and +then sends an event-stream exception frame and ends the stream, as Bedrock does for a mid-call +model error. ``streams()`` groups the inbound events per opened stream. """ from __future__ import annotations @@ -26,22 +32,24 @@ import uuid from collections.abc import Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass -from itertools import chain +from itertools import chain, count from pathlib import Path from queue import SimpleQueue -from typing import Final, Literal +from typing import Final, Literal, TypeVar import h2.config import h2.connection import h2.events import h2.exceptions from integration._support.async_server import Handler, LoopbackServer -from integration._support.process import _free_port +from integration._support.ports import allocate_port from integration._support.tls import server_context from integration._support.upstream import aws_event_stream_frame from pydantic import JsonValue, TypeAdapter BedrockEvent = Mapping[str, JsonValue] +OpeningEnd = Literal["system_prompt", "user_input"] +Recorded = TypeVar("Recorded") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -177,6 +185,20 @@ class Turn: events: tuple[BedrockEvent, ...] +@dataclass(frozen=True, slots=True) +class Refusal: + after: OpeningEnd + status: int = 400 + error_type: str = "ValidationException" + message: str = "The provided model identifier is invalid." + + +@dataclass(frozen=True, slots=True) +class StreamFailure: + exception_type: str = "modelStreamErrorException" + message: str = "The system encountered an unexpected error during processing. Try your request again." + + def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: body: Final = event["event"] assert isinstance(body, Mapping) and len(body) == 1, event @@ -214,6 +236,17 @@ def _frame_payload(raw: bytes) -> bytes: return raw[12 + header_length : -4] +def _exception_frame(failure: StreamFailure) -> bytes: + return aws_event_stream_frame( + { + ":message-type": "exception", + ":exception-type": failure.exception_type, + ":content-type": "application/json", + }, + json.dumps({"message": failure.message}).encode(), + ) + + def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" inner_payload: Final = _frame_payload(chunk_payload) @@ -239,6 +272,7 @@ class _Inbound: self.tool_use_ids: Final[dict[str, str]] = {} self.tool_results_seen = 0 self.user_text_closed = False + self.system_text_closed = False def record(self, event: BedrockEvent) -> None: self.events.put(event) @@ -253,6 +287,8 @@ class _Inbound: role_type: Final = self.open_blocks.pop(content_name, ("", "")) if role_type == ("USER", "TEXT"): self.user_text_closed = True + elif role_type == ("SYSTEM", "TEXT"): + self.system_text_closed = True elif name == "toolResult": self.tool_results_seen += 1 @@ -283,10 +319,35 @@ class _Script: return turn.events +def _opening_arrived(after: OpeningEnd, inbound: _Inbound) -> bool: + match after: + case "system_prompt": + return inbound.system_text_closed + case "user_input": + return inbound.user_text_closed + + class _PeerState: - def __init__(self, turns: tuple[Turn, ...]) -> None: + def __init__(self, turns: tuple[Turn, ...], refusal: Refusal | None, failure: StreamFailure | None) -> None: self.inbound: Final = _Inbound() self.script: Final = _Script(turns) + self.refusal: Final = refusal + self.failure: Final = failure + self.ordinals: Final = count() + self.opened: Final = SimpleQueue[int]() + self.stream_events: Final = SimpleQueue[tuple[int, BedrockEvent]]() + + +def _drain(queue: SimpleQueue[Recorded], log: list[Recorded]) -> tuple[Recorded, ...]: + while not queue.empty(): + log.append(queue.get()) # drained log, not a built value + return tuple(log) + + +@dataclass(frozen=True, slots=True) +class _Stream: + ordinal: int + inbound: _Inbound @dataclass(frozen=True, slots=True) @@ -294,11 +355,16 @@ class BedrockRealtimePeer: api_base: str _state: _PeerState _drained: list[BedrockEvent] + _opened: list[int] + _stream_events: list[tuple[int, BedrockEvent]] def received(self) -> tuple[BedrockEvent, ...]: - while not self._state.inbound.events.empty(): - self._drained.append(self._state.inbound.events.get()) # drained log, not a built value - return tuple(self._drained) + return _drain(self._state.inbound.events, self._drained) + + def streams(self) -> tuple[tuple[BedrockEvent, ...], ...]: + opened: Final = _drain(self._state.opened, self._opened) + recorded: Final = _drain(self._state.stream_events, self._stream_events) + return tuple(tuple(event for ordinal, event in recorded if ordinal == stream) for stream in opened) def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) @@ -344,6 +410,7 @@ class _H2Session: self._writer: Final = writer self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) self._buffers: Final[dict[int, bytes]] = {} + self._streams: Final[dict[int, _Stream]] = {} self._ended: set[int] = set() self._closed = False @@ -367,10 +434,14 @@ class _H2Session: def _apply(self, event: h2.events.Event) -> None: if isinstance(event, h2.events.RequestReceived): - self._h2.send_headers( - event.stream_id, - [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], - ) + ordinal: Final = next(self._state.ordinals) + self._streams[event.stream_id] = _Stream(ordinal, _Inbound()) + self._state.opened.put(ordinal) + if self._state.refusal is None: + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) elif isinstance(event, h2.events.DataReceived): self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data @@ -380,15 +451,44 @@ class _H2Session: decoded: Final = _decode_inbound(payload) if decoded is None: continue - self._state.inbound.record(decoded) - if event.stream_id not in self._ended: - for outbound in self._state.script.ready_events(self._state.inbound, decoded): - self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + self._record(event.stream_id, decoded) elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): self._close_stream(event.stream_id) elif isinstance(event, h2.events.ConnectionTerminated): self._writer.close() + def _record(self, stream_id: int, decoded: BedrockEvent) -> None: + stream: Final = self._streams[stream_id] + self._state.inbound.record(decoded) + stream.inbound.record(decoded) + self._state.stream_events.put((stream.ordinal, decoded)) + if stream_id in self._ended: + return + refusal: Final = self._state.refusal + if refusal is not None: + if _opening_arrived(refusal.after, stream.inbound): + self._refuse(stream_id, refusal) + return + outbound_events: Final = self._state.script.ready_events(self._state.inbound, decoded) + for outbound in outbound_events: + self._h2.send_data(stream_id, _outbound_chunk(outbound)) + failure: Final = self._state.failure + if outbound_events and failure is not None and self._state.script.exhausted: + self._h2.send_data(stream_id, _exception_frame(failure)) + self._end_stream(stream_id) + + def _refuse(self, stream_id: int, refusal: Refusal) -> None: + self._h2.send_headers( + stream_id, + [ + (":status", str(refusal.status)), + ("content-type", "application/json"), + ("x-amzn-errortype", refusal.error_type), + ], + ) + self._h2.send_data(stream_id, json.dumps({"message": refusal.message}).encode()) + self._end_stream(stream_id) + def _close_stream(self, stream_id: int) -> None: self._end_stream(stream_id) self._closed = True @@ -411,19 +511,27 @@ def _session_handler(state: _PeerState) -> Handler: @contextmanager -def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: +def bedrock_realtime_peer( + turns: tuple[Turn, ...], + cert: tuple[Path, Path], + *, + refusal: Refusal | None = None, + failure: StreamFailure | None = None, +) -> Iterator[BedrockRealtimePeer]: """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. SigV4 headers are ignored; the stream closes cleanly when the client ends it. """ - port: Final = _free_port() + port: Final = allocate_port() context: Final = server_context(cert[0], cert[1]) context.set_alpn_protocols(["h2"]) - state: Final = _PeerState(turns) + state: Final = _PeerState(turns, refusal, failure) server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) server.start() try: - yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + yield BedrockRealtimePeer( + api_base=f"https://localhost:{port}", _state=state, _drained=[], _opened=[], _stream_events=[] + ) finally: server.stop() 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/upstream.py b/tests/integration/_support/upstream.py index d2b67c5438e..186bc345fab 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -9,9 +9,10 @@ import struct import uuid import zlib from collections import deque -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from dataclasses import dataclass, field from queue import SimpleQueue +from types import MappingProxyType from typing import Final, cast import httpx @@ -42,7 +43,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route, WebSocketRoute -from starlette.websockets import WebSocket, WebSocketDisconnect +from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) INTERNAL_FIELDS: Final = frozenset( @@ -81,6 +82,25 @@ class Observation: api_key: str = "" +@dataclass(frozen=True, slots=True) +class _RealtimeConnection: + websocket: WebSocket + response: RealtimeResponse + session: dict[str, JsonValue] + scenario_id: str + authorization: str + api_key: str + events: Iterator[dict[str, JsonValue]] + typed_turn_detection: deque[JsonValue] + + def auto_answers(self) -> bool: + return ( + self.response.dialect == "xai" + and len(self.typed_turn_detection) > 0 + and _is_typed_server_vad(self.typed_turn_detection[-1]) + ) + + class InteractionState(BaseModel): """What the scripted Interactions API answers for one interaction id until a DELETE drops it.""" @@ -92,6 +112,12 @@ class InteractionState(BaseModel): delay_seconds: float = 0 +class _ReceivedFrame(BaseModel): + model_config = ConfigDict(frozen=True) + type: str + text: str | None = None + + class _ScenarioRegistration(BaseModel): scenario_id: str response: StoredResponse @@ -415,23 +441,47 @@ class Provider: ) -> None: await websocket.accept() session: Final = _realtime_session(response, scenario_id, websocket.query_params.get("model", "")) + frames: Final = _received_texts(websocket) + held: Final = await anext(frames, None) if response.created_after_first_frame else None for _ in range(response.created_repeats): await websocket.send_json({"type": response.created_event, "event_id": _event_id(), "session": session}) - event_index: Final = iter(response.events) - async for message in websocket.iter_json(): - payload: Final = JSON_OBJECT.validate_python(message) - self.observations.put(Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", api_key)) - if payload.get("type") == "session.update": - await websocket.send_json( - _realtime_update_reply(payload.get("session"), session, response.session_type) + connection: Final = _RealtimeConnection( + websocket, response, session, scenario_id, authorization, api_key, iter(response.events), deque(maxlen=1) + ) + if held is not None: + await self._realtime_frame(connection, held) + async for message in frames: + await self._realtime_frame(connection, message) + + async def _realtime_frame(self, connection: _RealtimeConnection, message: str) -> None: + websocket: Final = connection.websocket + response: Final = connection.response + payload: Final = JSON_OBJECT.validate_json(message) + self.observations.put( + Observation(websocket.url.path, connection.authorization, payload, "WEBSOCKET_FRAME", connection.api_key) + ) + frame_type: Final = payload.get("type") + if frame_type == "session.update": + reply: Final = _realtime_update_reply( + payload.get("session"), connection.session, response.session_type, response.dialect + ) + if reply["type"] == "session.updated": + connection.typed_turn_detection.extend( + filter(_declares_a_type, _declared_turn_detection(payload.get("session"))) ) - continue - if payload.get("type") not in _REALTIME_TRIGGERS: - continue - event: Final = next(event_index, None) - if event is None: - continue - await websocket.send_json(_rendered_realtime_event(event, scenario_id)) + await _reply_while_open(websocket, reply) + return + if frame_type == "response.cancel": + await _reply_while_open(websocket, _cancel_not_active(payload.get("event_id"), response.dialect)) + return + if frame_type not in _REALTIME_TRIGGERS: + return + event: Final = next(connection.events, None) + if event is None: + return + await _reply_while_open(websocket, _rendered_realtime_event(event, connection.scenario_id)) + if frame_type == "input_audio_buffer.commit" and connection.auto_answers(): + await _reply_while_open(websocket, _rendered_realtime_event(dict(_AUTO_ANSWER), connection.scenario_id)) async def muse_realtime(self, websocket: WebSocket) -> None: await websocket.accept() @@ -454,8 +504,23 @@ class Provider: await websocket.send_json(_rendered_realtime_event(pending.popleft(), scenario_id)) async def gemini_live(self, websocket: WebSocket) -> None: - authorization: Final = websocket.headers.get("authorization", "") - scenario_id: Final = websocket.headers.get("x-goog-user-project", "") + await self._gemini_live_session( + websocket, + websocket.headers.get("x-goog-user-project", ""), + websocket.headers.get("authorization", ""), + _accepts_every_setup, + ) + + async def gemini_ai_studio_live(self, websocket: WebSocket) -> None: + await self._gemini_live_session(websocket, websocket.query_params.get("key", ""), "", _ai_studio_setup_refusal) + + async def _gemini_live_session( + self, + websocket: WebSocket, + scenario_id: str, + authorization: str, + setup_refusal: Callable[[Mapping[str, JsonValue]], str | None], + ) -> None: self.observations.put( Observation( websocket.url.path, @@ -477,6 +542,9 @@ class Provider: Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", scenario_id) ) if "setup" in payload: + if (refusal := setup_refusal(payload)) is not None: + await websocket.close(code=_INVALID_ARGUMENT_CLOSE, reason=refusal[:_CLOSE_REASON_BYTES]) + return await websocket.send_json({"setupComplete": {}}) continue if not _GEMINI_LIVE_TRIGGERS.intersection(payload): @@ -587,6 +655,7 @@ class Provider: WebSocketRoute("/openai/realtime", self.realtime), WebSocketRoute("/v1/asr/realtime", self.muse_realtime), WebSocketRoute(GEMINI_LIVE_PATH, self.gemini_live), + WebSocketRoute(GEMINI_AI_STUDIO_LIVE_PATH, self.gemini_ai_studio_live), ] ) @@ -612,6 +681,19 @@ def _interaction_body(interaction_id: str, state: InteractionState) -> dict[str, _REALTIME_TRIGGERS: Final = frozenset({"response.create", "input_audio_buffer.commit"}) _GEMINI_LIVE_TRIGGERS: Final = frozenset({"clientContent", "realtimeInput"}) GEMINI_LIVE_PATH: Final = "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +GEMINI_AI_STUDIO_LIVE_PATH: Final = "/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" +_INVALID_ARGUMENT_CLOSE: Final = 1007 +_CLOSE_REASON_BYTES: Final = 123 +_AI_STUDIO_TUNING_FIELDS: Final = MappingProxyType( + {"prefixPaddingMs": "prefix_padding_ms", "silenceDurationMs": "silence_duration_ms"} +) +_AUTO_ANSWER: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": {"id": "resp_auto_$UNIQUE_ID", "object": "realtime.response", "status": "completed", "output": []}, + } +) _TRANSCRIPTION_UPDATE_REFUSED: Final = "Passing a realtime session update to a transcription session is not allowed." _REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." _NESTED_TURN_DETECTION_TYPE: Final = "session.audio.input.turn_detection.type" @@ -651,11 +733,11 @@ def _missing_turn_detection_type(session: Mapping[str, JsonValue]) -> str | None def _realtime_update_reply( - session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None + session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None, dialect: str ) -> dict[str, JsonValue]: if not isinstance(session, dict): return _realtime_error("invalid_value", "Invalid value for 'session': expected an object.", "session") - missing: Final = _missing_turn_detection_type(session) + missing: Final = None if dialect == "xai" else _missing_turn_detection_type(session) if missing is not None: return _realtime_error("missing_required_parameter", f"Missing required parameter: '{missing}'.", missing) declared: Final = session.get("type") @@ -666,6 +748,82 @@ def _realtime_update_reply( return {"type": "session.updated", "event_id": _event_id(), "session": {**created, **session}} +async def _received_texts(websocket: WebSocket) -> AsyncIterator[str]: + while (frame := _ReceivedFrame.model_validate(await websocket.receive())).type == "websocket.receive": + if frame.text is not None: + yield frame.text + + +async def _reply_while_open(websocket: WebSocket, frame: Mapping[str, JsonValue]) -> None: + if websocket.application_state is not WebSocketState.CONNECTED: + return + try: + await websocket.send_json(frame) + except WebSocketDisconnect: + return + + +def _cancel_not_active(cancel_event_id: JsonValue, dialect: str) -> dict[str, JsonValue]: + message: Final = "Cancellation failed: no active response found" + if dialect == "xai": + own_id: Final = _event_id() + return { + "type": "error", + "event_id": own_id, + "error": { + "message": message, + "type": "invalid_request_error", + "code": "invalid_request_error", + "event_id": own_id, + }, + } + return { + "type": "error", + "event_id": _event_id(), + "error": { + "type": "invalid_request_error", + "code": "response_cancel_not_active", + "message": message, + "param": None, + "event_id": cancel_event_id, + }, + } + + +def _declared_turn_detection(session: JsonValue) -> tuple[JsonValue, ...]: + if not isinstance(session, dict): + return () + audio: Final = session.get("audio") + audio_input: Final = audio.get("input") if isinstance(audio, dict) else None + if isinstance(audio_input, dict) and "turn_detection" in audio_input: + return (audio_input["turn_detection"],) + return (session["turn_detection"],) if "turn_detection" in session else () + + +def _declares_a_type(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and "type" in turn_detection + + +def _is_typed_server_vad(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and turn_detection.get("type") == "server_vad" + + +def _accepts_every_setup(_payload: Mapping[str, JsonValue]) -> str | None: + return None + + +def _ai_studio_setup_refusal(payload: Mapping[str, JsonValue]) -> str | None: + setup: Final = payload.get("setup") + config: Final = setup.get("realtimeInputConfig") if isinstance(setup, dict) else None + detection: Final = config.get("automaticActivityDetection") if isinstance(config, dict) else None + if not isinstance(detection, dict) or detection.get("disabled") is not True: + return None + tuning: Final = tuple(name for key, name in _AI_STUDIO_TUNING_FIELDS.items() if key in detection) + if not tuning: + return None + return f"automatic_activity_detection.disabled is true, but the following fields were also set: {', '.join(tuning)}" + + def _rendered_realtime_event(event: Mapping[str, JsonValue], scenario_id: str) -> dict[str, JsonValue]: return JSON_OBJECT.validate_json( json.dumps(event, separators=(",", ":")) 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/configuration/test_callback_delete.py b/tests/integration/configuration/test_callback_delete.py new file mode 100644 index 00000000000..d58dd252593 --- /dev/null +++ b/tests/integration/configuration/test_callback_delete.py @@ -0,0 +1,351 @@ +import json +import os +import uuid +from collections.abc import Generator, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, JsonValue, eventually, list_value, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server + +CALLBACK: Final = "generic_api" +SIBLING: Final = "langsmith" +LIST_TYPES: Final = MappingProxyType( + { + "success_callback": "success", + "failure_callback": "failure", + "callbacks": "success_and_failure", + } +) +SINK_ENVIRONMENT: Final = MappingProxyType( + { + "GENERIC_LOGGER_ENDPOINT": "http://127.0.0.1:9/unused", + "LANGSMITH_API_KEY": "synthetic-langsmith-key", + "LANGSMITH_BASE_URL": "http://127.0.0.1:9", + } +) +NOT_FOUND: Final = f"Callback '{CALLBACK}' not found in active configuration" +WORKER_SYNC_SECONDS: Final = 40 + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + database_url: str + + +@contextmanager +def _rig(gateway: Gateway, tmp_path: Path, file_settings: Mapping[str, JsonValue]) -> Generator[Rig]: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = tmp_path / f"callback_delete_{uuid.uuid4().hex}.yaml" + path.write_text( + yaml.safe_dump({**config, "litellm_settings": {**object_value(config["litellm_settings"]), **file_settings}}) + ) + with ( + scratch_database() as database_url, + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": database_url, **SINK_ENVIRONMENT}, + config=path, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as proxy, + ): + yield Rig(proxy, database_url) + + +@pytest.fixture +def rig(gateway: Gateway, tmp_path: Path) -> Iterator[Rig]: + with _rig(gateway, tmp_path, {}) as created: + yield created + + +def _configured(proxy: Gateway) -> tuple[tuple[str, str], ...]: + rows: Final = proxy.get("/get/config/callbacks")["callbacks"] + assert isinstance(rows, list) + return tuple( + sorted( + (str(row["name"]), str(row["type"])) + for row in (object_value(entry) for entry in rows) + if not row.get("read_only") + ) + ) + + +def _listed(proxy: Gateway, expected: tuple[tuple[str, str], ...]) -> None: + eventually(lambda: _configured(proxy), lambda rows: rows == expected, seconds=WORKER_SYNC_SECONDS) + + +def _stored_settings(database_url: str) -> Mapping[str, JsonValue]: + rows: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("litellm_settings",), + database_url=database_url, + ) + assert len(rows) == 1, rows + value: Final = rows[0]["param_value"] + return object_value(json.loads(value) if isinstance(value, str) else value) + + +def _configure( + proxy: Gateway, settings: Mapping[str, JsonValue], environment: Mapping[str, JsonValue] | None = None +) -> None: + body: Final[Mapping[str, JsonValue]] = { + "litellm_settings": dict(settings), + "environment_variables": None if environment is None else dict(environment), + } + assert proxy.post("/config/update", body) == {"message": "Config updated successfully"} + + +def _delete(proxy: Gateway, name: str) -> tuple[int, Mapping[str, JsonValue]]: + response: Final = proxy.request("POST", "/config/callback/delete", {"callback_name": name}) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +@pytest.mark.parametrize("list_key", tuple(LIST_TYPES)) +def test_callback_listed_under_any_settings_list_is_deleted_from_the_listing_and_the_store( + rig: Rig, list_key: str +) -> None: + _configure(rig.proxy, {list_key: [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, LIST_TYPES[list_key]),)) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (CALLBACK, []) + assert _stored_settings(rig.database_url)[list_key] == [] + _listed(rig.proxy, ()) + + +def test_callback_in_both_success_and_failure_lists_is_removed_from_both_in_one_delete(rig: Rig) -> None: + _configure(rig.proxy, {"success_callback": [CALLBACK], "failure_callback": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + stored: Final = _stored_settings(rig.database_url) + assert (stored["success_callback"], stored["failure_callback"]) == ([], []) + _listed(rig.proxy, ()) + + +def test_delete_keeps_other_callbacks_and_reports_what_is_left_in_the_list_it_edited(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING, CALLBACK], "failure_callback": [SIBLING]}) + _listed(rig.proxy, ((CALLBACK, "success_and_failure"), (SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [SIBLING] + stored: Final = _stored_settings(rig.database_url) + assert (stored["callbacks"], stored["failure_callback"]) == ([SIBLING], [SIBLING]) + _listed(rig.proxy, ((SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + +def test_delete_matches_the_callback_name_case_insensitively(rig: Rig) -> None: + _configure(rig.proxy, {"failure_callback": [CALLBACK.upper()], "callbacks": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK.upper(), "failure"), (CALLBACK, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK.upper()) + + assert status == 200, body + assert body["removed_callback"] == CALLBACK + _listed(rig.proxy, ()) + + +def test_delete_of_an_unconfigured_callback_is_a_404_and_changes_nothing(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING]}) + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + before: Final = _stored_settings(rig.database_url) + + status, body = _delete(rig.proxy, CALLBACK) + + assert (status, body) == (404, {"detail": {"error": NOT_FOUND}}) + assert _stored_settings(rig.database_url) == before + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + + +@pytest.mark.parametrize( + ("file_key", "database_key"), (("failure_callback", "success_callback"), ("success_callback", "failure_callback")) +) +def test_delete_removes_the_database_copy_and_keeps_the_copy_the_config_file_sets( + gateway: Gateway, tmp_path: Path, file_key: str, database_key: str +) -> None: + with _rig(gateway, tmp_path, {file_key: [CALLBACK]}) as file_owned: + _configure(file_owned.proxy, {database_key: [CALLBACK]}) + _listed(file_owned.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [] + assert _stored_settings(file_owned.database_url)[database_key] == [] + _listed(file_owned.proxy, ((CALLBACK, LIST_TYPES[file_key]),)) + + +def test_delete_of_a_callback_only_the_config_file_lists_is_refused_and_names_the_file( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, {"callbacks": [CALLBACK]}) as file_owned: + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 400, body + detail: Final = object_value(body["detail"]) + assert (detail["keys"], detail["section"]) == (["callbacks"], "litellm_settings") + assert "set in the config file" in str(detail["error"]) + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + +def _completion(request: Request) -> Reply: + message: Final = object_value(list_value(JSON_OBJECT.validate_json(request.body)["messages"])[0]) + if str(message["content"]).endswith("-failure"): + return Reply( + status=400, + body=json.dumps( + {"error": {"type": "invalid_request_error", "code": "synthetic", "message": "synthetic failure"}} + ).encode(), + ) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +@dataclass(slots=True) +class Sink: + wire: Wire + bodies: list[bytes] # mutable-ok: drain() consumes the queue, later reads keep earlier batches + + def received(self, markers: tuple[str, ...]) -> bool: + self.bodies.extend(request.body for request in self.wire.drain()) + return all(any(marker.encode() in body for body in self.bodies) for marker in markers) + + +@dataclass(frozen=True, slots=True) +class Pair: + first: Gateway + peer: Gateway + database_url: str + provider: Wire + generic: Sink + langsmith: Sink + + def running(self, name: str) -> tuple[bool, bool]: + return (_runs(self.first, name), _runs(self.peer, name)) + + +def _runs(proxy: Gateway, name: str) -> bool: + active: Final = list_value(proxy.get("/active/callbacks")["all_litellm_callbacks"]) + return any(name in str(entry).lower() for entry in active) + + +@contextmanager +def _pair(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[Pair]: + with ( + scratch_database() as database_url, + wire_server(_completion) as provider, + wire_server(lambda _: Reply()) as generic, + wire_server(lambda _: Reply()) as langsmith, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + environment: Final = { + "DATABASE_URL": database_url, + "DEFAULT_FLUSH_INTERVAL_SECONDS": "1", + "LITELLM_STORE_AUDIT_LOGS": "True", + } + removed: Final = ( + "DATABASE_URL_READ_REPLICA", + "GENERIC_LOGGER_ENDPOINT", + "GENERIC_LOGGER_HEADERS", + "LANGSMITH_API_KEY", + "LANGSMITH_BASE_URL", + ) + with ( + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as first, + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as peer, + ): + yield Pair(first, peer, database_url, provider, Sink(generic, []), Sink(langsmith, [])) + + +def _send_on_both(pair: Pair, model: str, key: str, marker: str, outcome: str) -> tuple[str, str]: + markers: Final = (f"{marker}-first-{outcome}", f"{marker}-peer-{outcome}") + for proxy, text in zip((pair.first, pair.peer), markers, strict=True): + response = proxy.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": text}]}, key=key + ) + assert response.status_code == (400 if outcome == "failure" else 200), response.text + return markers + + +@pytest.mark.parametrize(("list_key", "outcome"), (("failure_callback", "failure"), ("callbacks", "success"))) +def test_deleted_callback_stops_receiving_events_on_both_processes( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, list_key: str, outcome: str +) -> None: + assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: audit logs are an enterprise feature" + marker: Final = "callback-delete-" + uuid.uuid4().hex + with _pair(gateway, tmp_path, monkeypatch) as pair, pair.first.scenario() as scenario: + _configure( + pair.first, + {list_key: [SIBLING, CALLBACK]}, + { + "GENERIC_LOGGER_ENDPOINT": pair.generic.wire.url, + "LANGSMITH_API_KEY": "langsmith-key-" + marker, + "LANGSMITH_BASE_URL": pair.langsmith.wire.url, + }, + ) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + model: Final = scenario.model(api_base=pair.provider.url + "/v1") + key: Final = scenario.key(models=[model]) + before: Final = _send_on_both(pair, model, key, marker + "-before", outcome) + eventually(lambda: pair.generic.received(before), bool, seconds=25) + eventually(lambda: pair.langsmith.received(before), bool, seconds=25) + + status, body = _delete(pair.first, SIBLING) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (SIBLING, [CALLBACK]) + assert _stored_settings(pair.database_url)[list_key] == [CALLBACK] + audit: Final = eventually( + lambda: read_rows( + 'SELECT before_value, updated_values FROM "LiteLLM_AuditLog" WHERE object_id = %s AND action = %s', + ("litellm_settings", "deleted"), + database_url=pair.database_url, + ), + bool, + seconds=10, + ) + assert len(audit) == 1, audit + before_value: Final = list_value(object_value(audit[0]["before_value"])[list_key]) + assert sorted(str(name) for name in before_value) == sorted([SIBLING, CALLBACK]), audit + assert audit[0]["updated_values"] == {list_key: [CALLBACK]}, audit + eventually( + lambda: pair.running(SIBLING), lambda running: running == (False, False), seconds=WORKER_SYNC_SECONDS + ) + + after: Final = _send_on_both(pair, model, key, marker + "-after", outcome) + eventually(lambda: pair.generic.received(after), bool, seconds=25) + _configure(pair.first, {list_key: [CALLBACK, SIBLING]}) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + barrier: Final = _send_on_both(pair, model, key, marker + "-barrier", outcome) + eventually(lambda: pair.langsmith.received(barrier), bool, seconds=25) + + leaked: Final = tuple(text for text in after if any(text.encode() in body for body in pair.langsmith.bodies)) + assert leaked == (), leaked diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 3d74ac4e2f3..a0c81d8f06d 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4111,6 +4111,7 @@ CLAUDE_OPUS_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, ), @@ -4170,6 +4171,7 @@ CLAUDE_SONNET_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4232,6 +4234,7 @@ CLAUDE_SONNET_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4314,6 +4317,7 @@ CLAUDE_SONNET_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4566,6 +4570,7 @@ CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4628,6 +4633,7 @@ CLAUDE_HAIKU_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -4710,6 +4716,7 @@ CLAUDE_HAIKU_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5325,6 +5332,7 @@ CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5387,6 +5395,7 @@ CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5469,6 +5478,7 @@ CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5720,6 +5730,7 @@ CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( "output_tokens": 1, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -5802,6 +5813,7 @@ CLAUDE_OPUS_5_5_CACHE_READ: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6141,6 +6153,7 @@ CLAUDE_SONNET_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6294,6 +6307,7 @@ CLAUDE_HAIKU_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 3480}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6593,6 +6607,7 @@ CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -6746,6 +6761,7 @@ CLAUDE_OPUS_5_5_THINKING: Final = CostTrackingTestCase( "output_tokens_details": {"thinking_tokens": 3480}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 14c9647dc08..1360f65b802 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -2,7 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast from integration.cost_calculation.cost_tracking_case import ( with_response_usage, @@ -23,6 +23,30 @@ from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParit from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +def _response_for_model( + response: JsonResponse, model: str, *, service_tier: str | None = None +) -> JsonResponse: + return response.model_copy( + update={ + "body": { + **response.body, + "model": model, + **({"service_tier": service_tier} if service_tier is not None else {}), + } + } + ) + + +def _stream_for_model(response: SseResponse, model: str) -> SseResponse: + return response.model_copy( + update={ + "frames": tuple( + frame.replace("gpt-6-luna-2026-09-22", model).replace("gpt-6-luna", model) for frame in response.frames + ) + } + ) + + # azure/gpt-5.4-mini AZURE_GPT_5_4_MINI_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-5.4-mini-input_text", @@ -3604,173 +3628,7 @@ AZURE_GPT_6_LUNA_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6-luna-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-luna", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "895f84528ea4 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 895f84528ea4", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-luna-2026-09-22", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.0002535, input_cost=0.0001196, output_cost=0.0001339, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6-luna-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-luna", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "705c026752fc summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 705c026752fc", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-luna-2026-09-22", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.000507, input_cost=0.0002392, output_cost=0.0002678, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_ASTRA_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-6-astra-input_text", @@ -3983,173 +3841,7 @@ AZURE_GPT_6_ASTRA_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6-astra-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-astra", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6ff9027e3f7f summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 6ff9027e3f7f", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-astra-2026-09-03", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.024375, input_cost=0.0115, output_cost=0.012875, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6-astra-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6-astra", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "9a0c038a2ad8 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 9a0c038a2ad8", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6-astra-2026-09-03", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.04875, input_cost=0.023, output_cost=0.02575, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_1_SOL_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-6.1-sol-input_text", @@ -4362,173 +4054,7 @@ AZURE_GPT_6_1_SOL_STREAM: Final = CostTrackingTestCase( ), ) -AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="azure-gpt-6.1-sol-service_tier_flex", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6.1-sol", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "960e0935140f summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "flex", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 960e0935140f", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6.1-sol-2026-09-29", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "flex", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.00468, input_cost=0.002208, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 - ), -) -AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="azure-gpt-6.1-sol-service_tier_priority", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="azure/gpt-6.1-sol", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6b36e2dc1f33 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "service_tier": "priority", - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "choices": [ - { - "content_filter_results": _GPT_6_CONTENT_FILTER_RESULTS, - "finish_reason": "stop", - "index": 0, - "logprobs": None, - "message": { - "annotations": [], - "content": "scripted answer 6b36e2dc1f33", - "refusal": None, - "role": "assistant", - }, - } - ], - "created": 1791852029, - "id": "chatcmpl-$REQUEST_ID", - "model": "gpt-6.1-sol-2026-09-29", - "object": "chat.completion", - "prompt_filter_results": _GPT_6_PROMPT_FILTER_RESULTS, - "routing": {"serving_pipereplica": "d20260922011332-060efdec-prefill-r0-dp0-prefill"}, - "service_tier": "default", - "system_fingerprint": None, - "usage": { - "completion_tokens": 412, - "completion_tokens_details": { - "accepted_prediction_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - "rejected_prediction_tokens": 0, - }, - "latency_checkpoint": { - "engine_tbt_ms": 4, - "engine_ttft_ms": 131, - "engine_ttlt_ms": 157, - "pre_inference_ms": 55, - "service_tbt_ms": 5, - "service_ttft_ms": 1034, - "service_ttlt_ms": 1058, - "user_visible_ttft_ms": 979, - }, - "prompt_tokens": 1840, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - "total_tokens": 2252, - }, - }, - ), - expected=ExactExpected( - spend=0.00936, input_cost=0.004416, output_cost=0.004944, prompt_tokens=1840, completion_tokens=412 - ), -) AZURE_GPT_6_SOL_REASONING: Final = CostTrackingTestCase( name="azure/gpt-6-sol-reasoning", @@ -5981,31 +5507,7 @@ AZURE_GPT_6_LUNA_TIERED_AT_272K_BOUNDARY: Final = with_usage( ), ) -AZURE_GPT_6_LUNA_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX, - "azure-gpt-6-luna-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected( - spend=0.0392678, input_cost=0.039, output_cost=0.0002678, prompt_tokens=300000, completion_tokens=412 - ), -) -AZURE_GPT_6_LUNA_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY, - "azure-gpt-6-luna-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected( - spend=0.0785356, input_cost=0.078, output_cost=0.0005356, prompt_tokens=300000, completion_tokens=412 - ), -) AZURE_GPT_6_ASTRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( AZURE_GPT_6_ASTRA_INPUT_TEXT, @@ -6040,27 +5542,7 @@ AZURE_GPT_6_ASTRA_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=3.42575, input_cost=3.4, output_cost=0.02575, prompt_tokens=272000, completion_tokens=412), ) -AZURE_GPT_6_ASTRA_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX, - "azure-gpt-6-astra-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=3.77575, input_cost=3.75, output_cost=0.02575, prompt_tokens=300000, completion_tokens=412), -) -AZURE_GPT_6_ASTRA_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY, - "azure-gpt-6-astra-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=7.5515, input_cost=7.5, output_cost=0.0515, prompt_tokens=300000, completion_tokens=412), -) AZURE_GPT_6_1_SOL_TIERED_INPUT_ABOVE_272K: Final = with_usage( AZURE_GPT_6_1_SOL_INPUT_TEXT, @@ -6095,27 +5577,7 @@ AZURE_GPT_6_1_SOL_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=0.657744, input_cost=0.6528, output_cost=0.004944, prompt_tokens=272000, completion_tokens=412), ) -AZURE_GPT_6_1_SOL_TIERED_FLEX_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX, - "azure-gpt-6.1-sol-tiered_flex_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=0.724944, input_cost=0.72, output_cost=0.004944, prompt_tokens=300000, completion_tokens=412), -) -AZURE_GPT_6_1_SOL_TIERED_PRIORITY_ABOVE_272K: Final = with_usage( - AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY, - "azure-gpt-6.1-sol-tiered_priority_above_272k", - { - "prompt_tokens": 300000, - "total_tokens": 300412, - "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, - }, - ExactExpected(spend=1.449888, input_cost=1.44, output_cost=0.009888, prompt_tokens=300000, completion_tokens=412), -) AZURE_GPT_6_1_SOL_CACHE_WRITE_5M: Final = with_response_usage( AZURE_GPT_6_1_SOL_CACHE_READ, @@ -6174,6 +5636,154 @@ AZURE_GPT_6_LUNA_CACHE_WRITE_5M: Final = with_response_usage( "content_filter_result": {"error": {"code": "content_filter_error", "message": "The contents are not filtered"}} }, ) + +AZURE_GPT_5_6_TERRA_INPUT_TEXT: Final = AZURE_GPT_6_LUNA_INPUT_TEXT.model_copy( + update={ + "name": "azure-gpt-5.6-terra-input_text", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_INPUT_TEXT.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_STREAM: Final = AZURE_GPT_6_LUNA_STREAM.model_copy( + update={ + "name": "azure-gpt-5.6-terra-stream", + "model": "azure/gpt-5.6-terra", + "response": _stream_for_model( + cast(SseResponse, AZURE_GPT_6_LUNA_STREAM.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = AZURE_GPT_5_6_SERVICE_TIER_FLEX.model_copy( + update={ + "name": "azure-gpt-5.6-terra-service_tier_flex", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_5_6_SERVICE_TIER_FLEX.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0047432, + input_cost=0.002024, + output_cost=0.0027192, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = AZURE_GPT_5_6_SERVICE_TIER_PRIORITY.model_copy( + update={ + "name": "azure-gpt-5.6-terra-service_tier_priority", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_5_6_SERVICE_TIER_PRIORITY.response), + "gpt-5.6-terra-2026-07-09", + service_tier="default", + ), + "expected": ExactExpected( + spend=0.0189728, + input_cost=0.008096, + output_cost=0.0108768, + prompt_tokens=1840, + completion_tokens=412, + ), + "known_bug": "Azure GPT Terra priority service tier is billed at the base rate", + } +) +AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON: Final = AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "azure-gpt-5.6-terra-messages_bridged_json", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +AZURE_GPT_5_6_TERRA_CACHE_READ: Final = AZURE_GPT_6_LUNA_CACHE_READ.model_copy( + update={ + "name": "azure/gpt-5.6-terra-cache_read", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_CACHE_READ.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.00912736, + input_cost=0.00411136, + output_cost=0.005016, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M: Final = AZURE_GPT_6_LUNA_CACHE_WRITE_5M.model_copy( + update={ + "name": "azure/gpt-5.6-terra-cache_write_5m", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_CACHE_WRITE_5M.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.040216, + input_cost=0.0352, + output_cost=0.005016, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +AZURE_GPT_5_6_TERRA_REASONING: Final = AZURE_GPT_6_LUNA_REASONING.model_copy( + update={ + "name": "azure/gpt-5.6-terra-reasoning", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_REASONING.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.056056, + input_cost=0.002728, + output_cost=0.053328, + prompt_tokens=1240, + completion_tokens=4040, + ), + } +) +AZURE_GPT_5_6_TERRA_TOOL_CALL: Final = AZURE_GPT_6_LUNA_TOOL_CALL.model_copy( + update={ + "name": "azure/gpt-5.6-terra-tool_call", + "model": "azure/gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, AZURE_GPT_6_LUNA_TOOL_CALL.response), "gpt-5.6-terra-2026-07-09" + ), + "expected": ExactExpected( + spend=0.0094864, + input_cost=0.004048, + output_cost=0.0054384, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) AZURE_GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( AZURE_GPT_6_ASTRA_CACHE_READ, "azure/gpt-6-astra-cache_write_5m", @@ -6216,18 +5826,12 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_LUNA_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_LUNA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_LUNA_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_LUNA_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_LUNA_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_ASTRA_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_ASTRA_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_ASTRA_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_INPUT_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_6_1_SOL_TIERED_AT_272K_BOUNDARY, - AZURE_GPT_6_1_SOL_TIERED_FLEX_ABOVE_272K, - AZURE_GPT_6_1_SOL_TIERED_PRIORITY_ABOVE_272K, AZURE_GPT_5_4_MINI_INPUT_TEXT, AZURE_GPT_5_4_MINI_CACHE_READ, AZURE_GPT_5_4_MINI_AUDIO_INPUT, @@ -6276,16 +5880,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_SOL_SERVICE_TIER_PRIORITY, AZURE_GPT_6_LUNA_INPUT_TEXT, AZURE_GPT_6_LUNA_STREAM, - AZURE_GPT_6_LUNA_SERVICE_TIER_FLEX, - AZURE_GPT_6_LUNA_SERVICE_TIER_PRIORITY, AZURE_GPT_6_ASTRA_INPUT_TEXT, AZURE_GPT_6_ASTRA_STREAM, - AZURE_GPT_6_ASTRA_SERVICE_TIER_FLEX, - AZURE_GPT_6_ASTRA_SERVICE_TIER_PRIORITY, AZURE_GPT_6_1_SOL_INPUT_TEXT, AZURE_GPT_6_1_SOL_STREAM, - AZURE_GPT_6_1_SOL_SERVICE_TIER_FLEX, - AZURE_GPT_6_1_SOL_SERVICE_TIER_PRIORITY, AZURE_GPT_6_SOL_REASONING, AZURE_GPT_6_SOL_STREAM_NO_USAGE, AZURE_GPT_6_LUNA_REASONING, @@ -6304,6 +5902,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_1_SOL_CACHE_WRITE_5M, AZURE_GPT_6_LUNA_CACHE_WRITE_5M, AZURE_GPT_6_ASTRA_CACHE_WRITE_5M, + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + AZURE_GPT_5_6_TERRA_STREAM, + AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX, + AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, + AZURE_GPT_5_6_TERRA_CACHE_READ, + AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M, + AZURE_GPT_5_6_TERRA_REASONING, + AZURE_GPT_5_6_TERRA_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -6316,6 +5923,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=AZURE_GPT_6_LUNA_INPUT_TEXT, streamed=AZURE_GPT_6_LUNA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), + StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_INPUT_TEXT, streamed=AZURE_GPT_5_6_TERRA_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -6327,4 +5935,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=AZURE_GPT_6_LUNA_INPUT_TEXT, bridged=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=AZURE_GPT_6_ASTRA_INPUT_TEXT, bridged=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=AZURE_GPT_6_1_SOL_INPUT_TEXT, bridged=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase( + native=AZURE_GPT_5_6_TERRA_INPUT_TEXT, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON + ), ) diff --git a/tests/integration/cost_calculation/cases/bedrock_mantle.py b/tests/integration/cost_calculation/cases/bedrock_mantle.py index 714c6f2811f..2fa940d66f4 100644 --- a/tests/integration/cost_calculation/cases/bedrock_mantle.py +++ b/tests/integration/cost_calculation/cases/bedrock_mantle.py @@ -2,9 +2,23 @@ from __future__ import annotations +import json +from collections.abc import Mapping from typing import Final from integration.cost_calculation.cases.openai import with_usage +from integration.cost_calculation.cases.anthropic import ( + CLAUDE_OPUS_5_5_CACHE_READ, + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, + CLAUDE_OPUS_5_5_STREAM, + CLAUDE_OPUS_5_5_THINKING, + CLAUDE_OPUS_5_5_TOOL_CALL, +) +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -13,6 +27,7 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from pydantic import JsonValue BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON: Final = CostTrackingTestCase( @@ -503,6 +518,420 @@ BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_AT_272K_BOUNDARY: Final = with_usage( ExactExpected(spend=0.602932, input_cost=0.5984, output_cost=0.004532, prompt_tokens=272000, completion_tokens=412), ) +_MANTLE_GPT_RESPONSES_REQUEST: Final = { + "model": "$MODEL", + "input": "cost tracking fixture", + "max_output_tokens": 64, + "stream": False, +} +_MANTLE_GPT_REASONING_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "max_output_tokens": 128, + "reasoning": {"effort": "medium"}, +} +_MANTLE_GPT_STREAM_REQUEST: Final = {**_MANTLE_GPT_RESPONSES_REQUEST, "stream": True} +_MANTLE_GPT_CHAT_REQUEST: Final = { + "model": "$MODEL", + "messages": [{"role": "user", "content": "cost tracking fixture"}], + "max_completion_tokens": 64, +} + + +def _mantle_gpt_usage( + input_tokens: int, output_tokens: int, reasoning_tokens: int, cached_tokens: int = 0 +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": cached_tokens}, + "output_tokens": output_tokens, + "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, + "total_tokens": input_tokens + output_tokens, + } + + +def _mantle_gpt_json_case( + base: CostTrackingTestCase, + provider_model: str, + shape: str, + usage: Mapping[str, JsonValue], + expected: ExactExpected, + request: Mapping[str, JsonValue], +) -> CostTrackingTestCase: + case: Final = with_usage(base, f"bedrock_mantle-{provider_model}-{shape}", usage, expected) + if not isinstance(case.response, JsonResponse): + raise TypeError(f"{case.name}: expected a JSON response fixture") + response: Final = case.response.model_copy(update={"body": {**case.response.body, "model": provider_model}}) + return case.model_copy( + update={"model": f"bedrock_mantle/{provider_model}", "request": dict(request), "response": response} + ) + + +def _mantle_gpt_stream_case( + provider_model: str, usage: Mapping[str, JsonValue], expected: ExactExpected +) -> CostTrackingTestCase: + base: Final = BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM + if not isinstance(base.response, SseResponse): + raise TypeError(f"{base.name}: expected an SSE response fixture") + old_usage: Final = json.dumps(_mantle_gpt_usage(1840, 412, 0)) + new_usage: Final = json.dumps(dict(usage)) + frames: Final = tuple( + frame.replace("openai.gpt-6-sol", provider_model).replace(old_usage, new_usage) + for frame in base.response.frames + ) + response: Final = base.response.model_copy(update={"frames": frames}) + return base.model_copy( + update={ + "name": f"bedrock_mantle-{provider_model}-responses_stream", + "model": f"bedrock_mantle/{provider_model}", + "request": dict(_MANTLE_GPT_STREAM_REQUEST), + "response": response, + "expected": expected, + } + ) + + +def _mantle_gpt_cases( + provider_model: str, + response_reasoning_tokens: int, + response_expected: ExactExpected, + reasoning_usage: Mapping[str, JsonValue], + reasoning_expected: ExactExpected, + cache_reasoning_tokens: int, + cache_expected: ExactExpected, + tiered_input_expected: ExactExpected, + tiered_cache_expected: ExactExpected, + tiered_boundary_expected: ExactExpected, +) -> tuple[CostTrackingTestCase, ...]: + response_usage: Final = _mantle_gpt_usage(9, 64, response_reasoning_tokens) + response_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "responses_json", + response_usage, + response_expected, + _MANTLE_GPT_RESPONSES_REQUEST, + ) + stream_case: Final = _mantle_gpt_stream_case(provider_model, response_usage, response_expected) + reasoning_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_REASONING, + provider_model, + "responses_reasoning", + reasoning_usage, + reasoning_expected, + _MANTLE_GPT_REASONING_REQUEST, + ) + chat_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_CHAT_JSON, + provider_model, + "chat_json", + response_usage, + response_expected, + _MANTLE_GPT_CHAT_REQUEST, + ) + cache_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "responses_cache_read", + _mantle_gpt_usage(3144, 16, cache_reasoning_tokens, cached_tokens=3142), + cache_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 16}, + ) + tiered_input: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_input_above_272k", + _mantle_gpt_usage(300000, 412, 0), + tiered_input_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + tiered_cache: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_cache_read_above_272k", + _mantle_gpt_usage(300000, 412, 0, cached_tokens=200000), + tiered_cache_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + tiered_boundary: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tiered_at_272k_boundary", + _mantle_gpt_usage(272000, 412, 0), + tiered_boundary_expected, + {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, + ) + return (response_case, stream_case, reasoning_case, chat_case, cache_case, tiered_input, tiered_cache, tiered_boundary) + + +def _mantle_gpt_bridge(case: CostTrackingTestCase) -> BridgeCostParityTestCase: + bridged: Final = case.model_copy( + update={ + "name": case.name.removesuffix("responses_json") + "chat_bridged_json", + "endpoint": "/v1/chat/completions", + "request": dict(_MANTLE_GPT_CHAT_REQUEST), + } + ) + return BridgeCostParityTestCase(native=case, bridged=bridged) + + +_MANTLE_GPT_61_SOL_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6.1-sol", + 48, + ExactExpected(spend=0.00079618, input_cost=0.00002178, output_cost=0.0007744, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 5, 0), + ExactExpected(spend=0.00009922, input_cost=0.00003872, output_cost=0.0000605, prompt_tokens=16, completion_tokens=5), + 0, + ExactExpected(spend=0.000578622, input_cost=0.000385022, output_cost=0.0001936, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=1.4594778, input_cost=1.452, output_cost=0.0074778, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.5398778, input_cost=0.5324, output_cost=0.0074778, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.6632252, input_cost=0.65824, output_cost=0.0049852, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_6_LUNA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6-luna", + 64, + ExactExpected(spend=0.000043428, input_cost=0.000001188, output_cost=0.00004224, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 18, 11), + ExactExpected(spend=0.000013992, input_cost=0.000002112, output_cost=0.00001188, prompt_tokens=16, completion_tokens=18), + 16, + ExactExpected(spend=0.0000522984, input_cost=0.0000417384, output_cost=0.00001056, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=0.07960788, input_cost=0.0792, output_cost=0.00040788, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.03208788, input_cost=0.03168, output_cost=0.00040788, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.03617592, input_cost=0.035904, output_cost=0.00027192, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_6_ASTRA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-6-astra", + 29, + ExactExpected(spend=0.0047047, input_cost=0.0001287, output_cost=0.004576, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 5, 0), + ExactExpected(spend=0.0005863, input_cost=0.0002288, output_cost=0.0003575, prompt_tokens=16, completion_tokens=5), + 0, + ExactExpected(spend=0.00566566, input_cost=0.00452166, output_cost=0.001144, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=8.624187, input_cost=8.58, output_cost=0.044187, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=3.476187, input_cost=3.432, output_cost=0.044187, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=3.919058, input_cost=3.8896, output_cost=0.029458, prompt_tokens=272000, completion_tokens=412), +) +_MANTLE_GPT_56_TERRA_CASES: Final = _mantle_gpt_cases( + "openai.gpt-5.6-terra", + 33, + ExactExpected(spend=0.00121044, input_cost=0.00002772, output_cost=0.00118272, prompt_tokens=9, completion_tokens=64), + _mantle_gpt_usage(16, 17, 10), + ExactExpected(spend=0.00036344, input_cost=0.00004928, output_cost=0.00031416, prompt_tokens=16, completion_tokens=17), + 0, + ExactExpected(spend=0.001269576, input_cost=0.000973896, output_cost=0.00029568, prompt_tokens=3144, completion_tokens=16), + ExactExpected(spend=1.85942064, input_cost=1.848, output_cost=0.01142064, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.75062064, input_cost=0.7392, output_cost=0.01142064, prompt_tokens=300000, completion_tokens=412), + ExactExpected(spend=0.84537376, input_cost=0.83776, output_cost=0.00761376, prompt_tokens=272000, completion_tokens=412), +) + +_MANTLE_CLAUDE_OPUS_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" +_ADAPTIVE_USAGE: Final = { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 18, + "output_tokens_details": {"thinking_tokens": 6}, + "service_tier": "standard", +} +_ADAPTIVE_EXPECTED: Final = ExactExpected( + spend=0.0007524, input_cost=0.0001584, output_cost=0.000594, prompt_tokens=24, completion_tokens=18 +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( + CLAUDE_OPUS_5_5_INPUT_TEXT, + "bedrock_mantle-anthropic.claude-opus-5-5-input_text", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy( + update={ + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_INPUT_TEXT.request, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + } +) + +_MANTLE_CLAUDE_STREAM_RESPONSE: Final = SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-opus-5-5", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + }, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": ""}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "4"}}, + {"type": "content_block_stop", "index": 1}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": 24, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 18, + "output_tokens_details": {"thinking_tokens": 6}, + }, + }, + {"type": "message_stop"}, + ), +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM: Final = CLAUDE_OPUS_5_5_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-stream", + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_STREAM.request, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + "response": _MANTLE_CLAUDE_STREAM_RESPONSE, + "expected": _ADAPTIVE_EXPECTED, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = with_usage( + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + "bedrock_mantle-anthropic.claude-opus-5-5-messages_input_text", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy( + update={ + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.request, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + }, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-messages_stream", + "model": _MANTLE_CLAUDE_OPUS_MODEL, + "request": { + **CLAUDE_OPUS_5_5_MESSAGES_STREAM.request, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + }, + "response": _MANTLE_CLAUDE_STREAM_RESPONSE, + "expected": _ADAPTIVE_EXPECTED, + } +) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_READ, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_read", + { + "input_tokens": 22, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 4203, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 47, + "output_tokens_details": {"thinking_tokens": 35}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.00308319, input_cost=0.00153219, output_cost=0.001551, prompt_tokens=4225, completion_tokens=47 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_write_5m", + { + "input_tokens": 22, + "cache_creation_input_tokens": 4203, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 4203, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 44, + "output_tokens_details": {"thinking_tokens": 32}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.03627195, input_cost=0.03481995, output_cost=0.001452, prompt_tokens=4225, completion_tokens=44 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = with_usage( + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + "bedrock_mantle-anthropic.claude-opus-5-5-cache_write_1h", + { + "input_tokens": 22, + "cache_creation_input_tokens": 9003, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9003}, + "output_tokens": 12, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + }, + ExactExpected( + spend=0.1193808, input_cost=0.1189848, output_cost=0.000396, prompt_tokens=9025, completion_tokens=12 + ), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = with_usage( + CLAUDE_OPUS_5_5_THINKING, + "bedrock_mantle-anthropic.claude-opus-5-5-thinking", + _ADAPTIVE_USAGE, + _ADAPTIVE_EXPECTED, +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL: Final = with_usage( + CLAUDE_OPUS_5_5_TOOL_CALL, + "bedrock_mantle-anthropic.claude-opus-5-5-tool_call", + { + "input_tokens": 376, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 12, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + }, + ExactExpected(spend=0.0028776, input_cost=0.0024816, output_cost=0.000396, prompt_tokens=376, completion_tokens=12), +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-chat_bridged_json", + "response": BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.response, + } +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + _mantle_gpt_bridge(_MANTLE_GPT_61_SOL_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_6_LUNA_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_6_ASTRA_CASES[0]), + _mantle_gpt_bridge(_MANTLE_GPT_56_TERRA_CASES[0]), + BridgeCostParityTestCase( + native=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + bridged=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM, @@ -512,10 +941,35 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_INPUT_ABOVE_272K, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_CACHE_READ_ABOVE_272K, BEDROCK_MANTLE_OPENAI_GPT_6_SOL_TIERED_AT_272K_BOUNDARY, + *_MANTLE_GPT_61_SOL_CASES, + *_MANTLE_GPT_6_LUNA_CASES, + *_MANTLE_GPT_6_ASTRA_CASES, + *_MANTLE_GPT_56_TERRA_CASES, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, streamed=BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_STREAM ), + StreamParityTestCase(plain=_MANTLE_GPT_61_SOL_CASES[0], streamed=_MANTLE_GPT_61_SOL_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_6_LUNA_CASES[0], streamed=_MANTLE_GPT_6_LUNA_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_6_ASTRA_CASES[0], streamed=_MANTLE_GPT_6_ASTRA_CASES[1]), + StreamParityTestCase(plain=_MANTLE_GPT_56_TERRA_CASES[0], streamed=_MANTLE_GPT_56_TERRA_CASES[1]), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM + ), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index a4538b87db3..bd64f4aeb11 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -2,7 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( @@ -16,6 +16,32 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames + +def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: + match response: + case JsonResponse(): + return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + case SseResponse(): + return response.model_copy( + update={ + "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + } + ) + + +def _gemini_case( + sibling: CostTrackingTestCase, name: str, model: str, model_version: str, expected: ExactExpected +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "model": model, + "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), model_version), + "expected": expected, + } + ) + + # gemini/gemini-3.1-pro GEMINI_GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( name="gemini-gemini-3.1-pro-input_text", @@ -61,6 +87,7 @@ GEMINI_GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -115,6 +142,7 @@ GEMINI_GEMINI_3_1_PRO_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -360,6 +388,7 @@ GEMINI_GEMINI_3_1_PRO_REASONING: Final = CostTrackingTestCase( "thoughtsTokenCount": 3480, "totalTokenCount": 5280, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.1-pro", }, @@ -1574,254 +1603,9 @@ GEMINI_GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( ), ) -GEMINI_GEMINI_3_8_FLASH_AUDIO_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-audio_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "99e65f16c4b4 summarize the attached material in one line and name the city weather", - }, - { - "type": "input_audio", - "input_audio": { - "data": "UklGRmQGAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YUAGAAAAAA8I4A85F+EdpCNYKNkrCy7eLkwuWSwTKZUk/x59GEERgQl4AWb5hPER6kbjVd1t2LLUQdIu0X/RM9M81oTa6t9E5mPtD/UP/SQFEw2fFI0bqSHEJrgqZi27LqouNi1pKlkmJSH0GvUTXgxpBFP8WPS27KflYt8U2ujV/dJp0TjRbNL71NPY1d3b47nqOPIh+jUCOQrwER0ZjB8JJW0plCxoLtou5i2VK/cnKSNOHZUWLw9VB0T/OPdv7yToj+Hi20nX5tPT0SDR09Hm00nX4tuP4STob+8490T/VQcvD5UWTh0pI/cnlSvmLdouaC6ULG0pCSWMHx0Z8BE5CjUCIfo48rnq2+PV3dPY+9Rs0jjRadH90ujVFNpi36fltuxY9FP8aQReDPUT9BolIVkmaSo2Laouuy5mLbgqxCapIY0bnxQTDSQFD/0P9WPtRObq34TaPNYz03/RLtFB0rLUbdhV3UbjEeqE8Wb5eAGBCUERfRj/HpUkEylZLEwu3i4LLtkrWCikI+EdORfgDw8IAADx9yDwx+gf4lzcqNcn1PXRItG00afT7dZr2wHhg+e/7n/2iP6aBnwO7xW6HKsikydOK78t0i6BLs0sxCl8JRYgvBmdEvEK8QLc+u3yYetz5FfePNlI1ZrSRdFW0crSl9Wn2dveDOUL7KLzl/utA6gLShNZGp4g7CUYKgMtly7ILpQtBSstJysiJRxHFcgN3wXL/cf1EO7j5nTg99qT1mzTmNEm0RrSa9QJ2NfcsuJr6dHwq/i8AMgIkRDcF3EeHiS3KBosLS7gLi0uGiy3KB4kcR7cF5EQyAi8AKv40fBr6bLi19wJ2GvUGtIm0ZjRbNOT1vfadODj5hDux/XL/d8FyA1HFSUcKyItJwUrlC3ILpcuAy0YKuwlniBZGkoTqAutA5f7ovML7Azl296n2ZfVytJW0UXRmtJI1TzZV95z5GHr7fLc+vEC8QqdErwZFiB8JcQpzSyBLtIuvy1OK5MnqyK6HO8VfA6aBoj+f/a/7oPnAeFr2+3Wp9O00SLR9dEn1KjXXNwf4sfoIPDx9wAADwjgDzkX4R2kI1go2SsLLt4uTC5ZLBMplST/Hn0YQRGBCXgBZvmE8RHqRuNV3W3YstRB0i7Rf9Ez0zzWhNrq30TmY+0P9Q/9JAUTDZ8UjRupIcQmuCpmLbsuqi42LWkqWSYlIfQa9RNeDGkEU/xY9Lbsp+Vi3xTa6NX90mnRONFs0vvU09jV3dvjueo48iH6NQI5CvARHRmMHwklbSmULGgu2i7mLZUr9ycpI04dlRYvD1UHRP8492/vJOiP4eLbSdfm09PRINHT0ebTSdfi24/hJOhv7zj3RP9VBy8PlRZOHSkj9yeVK+Yt2i5oLpQsbSkJJYwfHRnwETkKNQIh+jjyuerb49Xd09j71GzSONFp0f3S6NUU2mLfp+W27Fj0U/xpBF4M9RP0GiUhWSZpKjYtqi67LmYtuCrEJqkhjRufFBMNJAUP/Q/1Y+1E5urfhNo81jPTf9Eu0UHSstRt2FXdRuMR6oTxZvl4AYEJQRF9GP8elSQTKVksTC7eLgsu2StYKKQj4R05F+APDwgAAPH3IPDH6B/iXNyo1yfU9dEi0bTRp9Pt1mvbAeGD57/uf/aI/poGfA7vFbocqyKTJ04rvy3SLoEuzSzEKXwlFiC8GZ0S8QrxAtz67fJh63PkV9482UjVmtJF0VbRytKX1afZ294M5QvsovOX+60DqAtKE1kaniDsJRgqAy2XLsgulC0FKy0nKyIlHEcVyA3fBcv9x/UQ7uPmdOD32pPWbNOY0SbRGtJr1AnY19yy4mvp0fCr+LwAyAiRENwXcR4eJLcoGiwtLuAuLS4aLLcoHiRxHtwXkRDICLwAq/jR8GvpsuLX3AnYa9Qa0ibRmNFs05PW99p04OPmEO7H9cv93wXIDUcVJRwrIi0nBSuULcguly4DLRgq7CWeIFkaShOoC60Dl/ui8wvsDOXb3qfZl9XK0lbRRdGa0kjVPNlX3nPkYevt8tz68QLxCp0SvBkWIHwlxCnNLIEu0i6/LU4rkyerIroc7xV8DpoGiP5/9r/ug+cB4Wvb7dan07TRItH10SfUqNdc3B/ix+gg8PH3", - "format": "wav", - }, - }, - ], - }, - ], - "stream": False, - "modalities": ["text"], - "allowed_openai_params": ["modalities"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 99e65f16c4b4"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1546, - "candidatesTokenCount": 210, - "totalTokenCount": 1756, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 96}, - {"modality": "AUDIO", "tokenCount": 1450}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.002128, input_cost=0.001498, output_cost=0.00063, prompt_tokens=1546, completion_tokens=210 - ), -) -GEMINI_GEMINI_3_8_FLASH_AUDIO_OUTPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-audio_output", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "6fc6e4823e02 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "modalities": ["text", "audio"], - "audio": {"voice": "alloy", "format": "pcm16"}, - "allowed_openai_params": ["modalities", "audio"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 6fc6e4823e02"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 220, - "candidatesTokenCount": 1300, - "totalTokenCount": 1520, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 220}], - "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 180}, - {"modality": "AUDIO", "tokenCount": 1120}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.00737, input_cost=0.00011, output_cost=0.00726, prompt_tokens=220, completion_tokens=1300 - ), -) -GEMINI_GEMINI_3_8_FLASH_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-image_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "ca377dd90846 summarize the attached material in one line and name the city weather", - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAMK0lEQVR4nAEgDN/zAM0HLNi+b59irEwJwoIG5+NVlKprNC9dCjpeSEL6tCj3YubiguXBZXx4w6lns2cR6zkGp8hgPXHUCeelTYe9wfcEQgJ6rx+pW3+GWJV430PkExZ66Nnc6zd2KDOBGnGnIwBzhiZIL2HGI3lifMEk1EYYPG5NnqGlpcz3LiFAwwS9/Goh5egSF1aINdSX+yErhrSeZWrPBkEWmgtZ9OYpQ58l2dRlT+yNSBn7QNa6ssjgEhxEGuYUp7jZKpIZrx7Oh1QAV1jeeB7zT4/0jccZhxGSWqLiJW9VRPJQuRZjnL/J8qO8F7vpSKdYNMGDc/cEMXKNBlAdehNaVHGe84Tdmm53hcOfr0ICjvEPuk0WzuaDIOumjneMSZ1+6qg8mAM8quAXAOyQPrjRNxDXsUwZZhYr07VMCinT0OP4yIEWDKvlaxGgReVKAJ1JpZx/Hlt+evT70yo3G94iWEhVavFwPnqJ87rKl0BT6+ohtOgz197MvB8QzMXpME/BwepPYkiRLpbBOAD37hU9bAWozdK3sPczg3okfSudzcFjAYtEIq5yw+1ZF9EYmBSexEP+IhnvUWC4BeDWZQiCjhF7/4syzu4C5kF9ETfrG+udK0232R+NA+uESn014bRJmPMfahYljDojL1UAbuaA0PuOGOzBBlCKXAkFNHIfvvZHQafMkl9qmqdFF4x3Em6WDuijSQXN6nFtM3UXbUGmmCF4Rcyl4YhifPspUXLdXZMIvPo9zAhTSlQFEi8m83sw5KxL0rWBzS9c4XAIAOaz3py4dTb7d9QaqDMLk0J87/15PZKvEaC6/hZ62sCtuFTywatjViHrBXTgOO5IJsiyYuzmaeQJKHmr11gnixR2rO7lqNEGs3shT+xK/lDUucFkigvA+a5C+itk/VV+1gDxc420UHpKhj31j0YnCZSFJebGz2/XSTyT6XfZhG4XNwZGIeUGCPKt0DX9lpB0RNN0yiPzQVJfa3LkZpRBN3RGgRpYc8WpHn5Q1wWpmnklpPHACv/N+kGz0Ki86g2GgvsAWloXy5pwfFtwZRYV9fMGU4Za35yfj4cddJuHfAyHSpYHVlGhNknUVQIBV9gOq7wwLZU3Pl5GJgSy4EK7f75iRVKD/B0ItpC0FBpwODhWP184ymnLgLGkK90WIVUY7hZtAEOu39BF2OsPDWrBGWN0esj6v3clX2z22gaLmrJHiwE4sXWUC8TLLtFp4uiS/VlbojLP9uiewb/vrTLBiFjYJ5rsFjuu/3XxEuiZ1QYyi9sfsFqPoiXjQjCA/jibX4aA1AAfp3GT1lykHtVMJmSxoW4Xvn3BXhXUdtWhIwP7NDO1HRP9UAlEUgKbd/iJBWS20DFBJQb2/UN/+CpSWi97Rta3+pe3H4kOr3qaV+g1Udgmupu6/cxmQqMP+DXd77Sy6a0AMxTVBR0DU4sVW/Vs2qPfnh3r+xlqt/3VIhyKQknN6xF5RIg4j7xsEpjsnKV/XhJNmN2sWekxom9lUCkuP3qgD25S7oCG6ZV3CrkUCm48s5hw+dUZANcGs4H6/PxIrSpkAC77CDPFF5hCvkfKWzerhufLBkq7AhZgeMKRnNboaP3m9qMh6+9Z3JEyaV8rfUacshIsMqwS/BI0uL9v9+fwcMQrbdwOstrkyWCPG61dXIAoEb9t2HnSdSlxy6FXa5+LhwCvCy1ANKkBHgVSx5hoE+LrBX47caskZapZ+MAsTFJhA3VxvHgKaWiuX4/vaG7TbOZdX7GRSzPz3yOeM4JeAuLqwOy6T0OBIKbnSm5b3A5+Y2j2cNYMiVmogx89QCIPRicA936Dj6rC2bDKBi8DmXw8dZvR170ELj4UjqD+VRopML3ww7ILyoVYi5P150ekt4QioS95PZdaHcPXSAD0GwVZe1t0K1qy2TGc7V2ySUwrZKwEnPRbLXAcl6praPJ8aVbkAJ1MPaI09ZDaZOT+npRQ4SFv1DK3halvT/cYVWNFxZy/HEwXatv4MtSF+5ymHEWqFC/kYwCJMTaYtzI7Ma9Q1rJ1WZtVBP36KFFdSj1G8Bw5b5ssozn/uHJxFO9g937ZtQBQvxvgA4h8rApfcpEHs+HfrYkWalhdEwiJ+fpmF/Um3ywbq7MF3kWROuUQa4H2rcVhq4WpdGuBte/A+QvHrGkq45kCcmvaWhAAslxCV5CWs7QlXig29URyTAgPh//Ji+IAJXC9fNNMdeirO7iPEEOemj9zZ8FMiAQAYKRF4o8E9glP+Jx+Vw9xU5oM40/363XW5T+Gd4xtwwwQoR3P38Wd0uURAeDvp2P53Wz6zxpyRmpc2iAwN5BujDYD2v+lpIn2AMUaEqKkrYP6sRheFZ0EY9Yt7r25gt9pITZVoPxxTMUDlHfdZgyMFfPLKLOtUk3gaiv88FD1Wd4AdI2pNnmYqAMKjqK35Ys3wVuBmgAbUPH6k4ae0ku7/qyRrkGHwRewnAAVZQgZavvBMFJ8cB4Ndb2bXEI2phv40TLtxqffTyNrTfbyrEdKKESwv/h/+pkLpi4bcaUZx0fBeRewm9oUozrsfuGNZcn0rKC8DdMUtqqXBWClU0bAHPTp2GD2c1FNwcwAFOXWyuK3ox0nFYLckD+pQW3ZJyKBKlCnanFXhVDRONHgupmm89QapXZ32FhAzncd/Xcy34ljiSee3UVHJtnq2YJJwQQVBNQKjuhoCv0YqzT/zVWuTbF5QtRm8I7cC5FQANj9TfJqwOxdaoY98K6RRZGzAezoRsF02Q3PwAowC6vYqUzHv85kEfb7BYqzoXGLm4fLxafznJZ+JhJdtsmV5qQiAxjnFQdrdTtL4KNZCXlvtdVYXxLGFN9os7WJoU2kLQB5RTU1rvAlYHfbTZZTTYFKXxRHAizqcSNldKkms7p4RYVAWlUXIGdWTb4kym3qAF4clO+3Njf9F3zxl1b73YfH8pXbnk9fIQnHRowK9vQFRcfD8ikeIUAmZrhe+3cOXZUAt7EeSjZgZFxWFhFll/ZC/Y93aYwMJjAhvbgcS/zZaR0ucmJ663gDu7xgXT3tgQrvb4daDFM10U5HZ/kt2r9m26me5oXdgom65hEkEgfTe03ZFzZnPz7n1fzuGVThCppMAAAyK2pIlNAb40xyGcWNko8VyJ3oghsue2leWHkJ6UpWCdZB13x+QcxkLOr5a8X0zx6N+VcXwNMfGGqlelK0sh1MrRizk4drf0prMW3Q264J3IVJTZ3tx1ULpBi8CnnmrwBeYvfkI5rZJ2W6cOv3bKMqoCpyoNo+gpCZcSVgB8hdzlfMfPl67dL9/IqNo2mBQBaVitoQ4y7KzqHnrBWSxLIyznP3v/VNlzgF7qcPaQiDYhyGAgLfihsZE5nQZLcedY4Apm3dhEOyMUpdzNNyViCP6ykfFn35Ss3JDUS7lfZV47bpxeatoe2dNzU5TGxY9GoVYvhKXmP0WjgOpwv91XPF/x1ty0JRgNGVMwnwxk4PpSpDgGRu/k4jMev9x1vlc3EdAIz+WBCWU23SZoIT3WfyVNpXi7azEQkSB7e1K+fS25uGiNScc4/Nn8BVsUqtwtGhM5nncB7tH0F/k0ngZeGmqEGNJw01v2UYpLQoxwqhCz7/X5it8liVJMxk2kQ3iCGYnwDwfHcl7gLJGiAqrjLBOlp17bPGYIGTlFw+TOCLsGS103TZvSbyptAazsca/TY3o5hZLqX7SBvnfL6s/SU7BVyQY4PpZwHLa6PI2w+uFmJZbBsaruCVPW+FL+zSfvPO560ABJLD6nKsReNBwoxShnNi1XPSmX+C7Ij/ik2nfzubk6IFlomrI0bs386ztTMDifWKNO0ZnC+Oau6ZgwE4PwF8AqKtGFd5yqG0aMFq9rvwPb8DZCXmUHN5Uc40cNG5HWNSAK1yFAkCcG2GPV/7JYEi3t9ZA8NBeh9Dmvsjv8jiaeknPCij14x6BmGOFzyVNqRcS8x517fM39S0cC6bzu8wbnh96fwQ9z/JzCEsqxV9KjuES+xv3rYkX+r3JxcK7p78EACPip85gRBoIU7cZgrwLepMArrpA9YvR8Icag3YILzWcgXVoVRTs6Xc2PsiKasZ98y5CBkmP9Ri3lwZ9acgvPJ1fw729yDzDl/0qEeBOf7ZYZ7dreaPeaGHIymfgKAwmwhnjvAGaqLWhgAAAABJRU5ErkJggg==", - "detail": "high", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer ca377dd90846"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 2116, - "candidatesTokenCount": 240, - "totalTokenCount": 2356, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 310}, - {"modality": "IMAGE", "tokenCount": 1806}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0018683, input_cost=0.0011483, output_cost=0.00072, prompt_tokens=2116, completion_tokens=240 - ), -) -GEMINI_GEMINI_3_8_FLASH_VIDEO_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-video_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "030071a5c80f summarize the attached material in one line and name the city weather", - }, - { - "type": "file", - "file": { - "file_data": "data:video/mp4;base64,AAAAGGZ0eXBpc29tAAACAGlzb21pc282AAACCG1kYXQNFBsiKTA3PkVMU1phaG92fYSLkpmgp661vMPK0djf5u30+wIJEBceJSwzOkFIT1ZdZGtyeYCHjpWco6qxuL/GzdTb4unw9/4FDBMaISgvNj1ES1JZYGdudXyDipGYn6attLvCydDX3uXs8/oBCA8WHSQrMjlAR05VXGNqcXh/ho2Um6KpsLe+xczT2uHo7/b9BAsSGSAnLjU8Q0pRWF9mbXR7gomQl56lrLO6wcjP1t3k6/L5AAcOFRwjKjE4P0ZNVFtiaXB3foWMk5qhqK+2vcTL0tng5+71/AMKERgfJi00O0JJUFdeZWxzeoGIj5adpKuyucDHztXc4+rx+P8GDRQbIikwNz5FTFNaYWhvdn2Ei5KZoKeutbzDytHY3+bt9PsCCRAXHiUsMzpBSE9WXWRrcnmAh46VnKOqsbi/xs3U2+Lp8Pf+BQwTGiEoLzY9REtSWWBnbnV8g4qRmJ+mrbS7wsnQ197l7PP6AQgPFh0kKzI5QEdOVVxjanF4f4aNlJuiqbC3vsXM09rh6O/2/QQLEhkgJy41PENKUVhfZm10e4KJkJeepayzusHIz9bd5Ovy+QAHDhUcIyoxOD9GTVRbYmlwd36FjJOaoaivtr3Ey9LZ4Ofu9fwDChEYHyYtNDtCSVBXXmVsc3qBiI+WnaSrsrnAx87V3OPq8fj/Bg==", - "format": "mp4", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 030071a5c80f"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 8060, - "candidatesTokenCount": 300, - "totalTokenCount": 8360, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 140}, - {"modality": "VIDEO", "tokenCount": 7920}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.005722, input_cost=0.004822, output_cost=0.0009, prompt_tokens=8060, completion_tokens=300 - ), -) GEMINI_GEMINI_3_8_FLASH_REASONING: Final = CostTrackingTestCase( name="gemini-gemini-3.8-flash-reasoning", @@ -2290,60 +2074,6 @@ GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( expected=RecountExpected(recount=RecountRates(input_cost_per_token=5e-07, output_cost_per_token=3e-06)), ) -GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-gemini-3.8-flash-stream_no_usage_image_input", - covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini/gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "5fca57bc9ae9 summarize the attached material in one line and name the city weather", - }, - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAMK0lEQVR4nAEgDN/zAM0HLNi+b59irEwJwoIG5+NVlKprNC9dCjpeSEL6tCj3YubiguXBZXx4w6lns2cR6zkGp8hgPXHUCeelTYe9wfcEQgJ6rx+pW3+GWJV430PkExZ66Nnc6zd2KDOBGnGnIwBzhiZIL2HGI3lifMEk1EYYPG5NnqGlpcz3LiFAwwS9/Goh5egSF1aINdSX+yErhrSeZWrPBkEWmgtZ9OYpQ58l2dRlT+yNSBn7QNa6ssjgEhxEGuYUp7jZKpIZrx7Oh1QAV1jeeB7zT4/0jccZhxGSWqLiJW9VRPJQuRZjnL/J8qO8F7vpSKdYNMGDc/cEMXKNBlAdehNaVHGe84Tdmm53hcOfr0ICjvEPuk0WzuaDIOumjneMSZ1+6qg8mAM8quAXAOyQPrjRNxDXsUwZZhYr07VMCinT0OP4yIEWDKvlaxGgReVKAJ1JpZx/Hlt+evT70yo3G94iWEhVavFwPnqJ87rKl0BT6+ohtOgz197MvB8QzMXpME/BwepPYkiRLpbBOAD37hU9bAWozdK3sPczg3okfSudzcFjAYtEIq5yw+1ZF9EYmBSexEP+IhnvUWC4BeDWZQiCjhF7/4syzu4C5kF9ETfrG+udK0232R+NA+uESn014bRJmPMfahYljDojL1UAbuaA0PuOGOzBBlCKXAkFNHIfvvZHQafMkl9qmqdFF4x3Em6WDuijSQXN6nFtM3UXbUGmmCF4Rcyl4YhifPspUXLdXZMIvPo9zAhTSlQFEi8m83sw5KxL0rWBzS9c4XAIAOaz3py4dTb7d9QaqDMLk0J87/15PZKvEaC6/hZ62sCtuFTywatjViHrBXTgOO5IJsiyYuzmaeQJKHmr11gnixR2rO7lqNEGs3shT+xK/lDUucFkigvA+a5C+itk/VV+1gDxc420UHpKhj31j0YnCZSFJebGz2/XSTyT6XfZhG4XNwZGIeUGCPKt0DX9lpB0RNN0yiPzQVJfa3LkZpRBN3RGgRpYc8WpHn5Q1wWpmnklpPHACv/N+kGz0Ki86g2GgvsAWloXy5pwfFtwZRYV9fMGU4Za35yfj4cddJuHfAyHSpYHVlGhNknUVQIBV9gOq7wwLZU3Pl5GJgSy4EK7f75iRVKD/B0ItpC0FBpwODhWP184ymnLgLGkK90WIVUY7hZtAEOu39BF2OsPDWrBGWN0esj6v3clX2z22gaLmrJHiwE4sXWUC8TLLtFp4uiS/VlbojLP9uiewb/vrTLBiFjYJ5rsFjuu/3XxEuiZ1QYyi9sfsFqPoiXjQjCA/jibX4aA1AAfp3GT1lykHtVMJmSxoW4Xvn3BXhXUdtWhIwP7NDO1HRP9UAlEUgKbd/iJBWS20DFBJQb2/UN/+CpSWi97Rta3+pe3H4kOr3qaV+g1Udgmupu6/cxmQqMP+DXd77Sy6a0AMxTVBR0DU4sVW/Vs2qPfnh3r+xlqt/3VIhyKQknN6xF5RIg4j7xsEpjsnKV/XhJNmN2sWekxom9lUCkuP3qgD25S7oCG6ZV3CrkUCm48s5hw+dUZANcGs4H6/PxIrSpkAC77CDPFF5hCvkfKWzerhufLBkq7AhZgeMKRnNboaP3m9qMh6+9Z3JEyaV8rfUacshIsMqwS/BI0uL9v9+fwcMQrbdwOstrkyWCPG61dXIAoEb9t2HnSdSlxy6FXa5+LhwCvCy1ANKkBHgVSx5hoE+LrBX47caskZapZ+MAsTFJhA3VxvHgKaWiuX4/vaG7TbOZdX7GRSzPz3yOeM4JeAuLqwOy6T0OBIKbnSm5b3A5+Y2j2cNYMiVmogx89QCIPRicA936Dj6rC2bDKBi8DmXw8dZvR170ELj4UjqD+VRopML3ww7ILyoVYi5P150ekt4QioS95PZdaHcPXSAD0GwVZe1t0K1qy2TGc7V2ySUwrZKwEnPRbLXAcl6praPJ8aVbkAJ1MPaI09ZDaZOT+npRQ4SFv1DK3halvT/cYVWNFxZy/HEwXatv4MtSF+5ymHEWqFC/kYwCJMTaYtzI7Ma9Q1rJ1WZtVBP36KFFdSj1G8Bw5b5ssozn/uHJxFO9g937ZtQBQvxvgA4h8rApfcpEHs+HfrYkWalhdEwiJ+fpmF/Um3ywbq7MF3kWROuUQa4H2rcVhq4WpdGuBte/A+QvHrGkq45kCcmvaWhAAslxCV5CWs7QlXig29URyTAgPh//Ji+IAJXC9fNNMdeirO7iPEEOemj9zZ8FMiAQAYKRF4o8E9glP+Jx+Vw9xU5oM40/363XW5T+Gd4xtwwwQoR3P38Wd0uURAeDvp2P53Wz6zxpyRmpc2iAwN5BujDYD2v+lpIn2AMUaEqKkrYP6sRheFZ0EY9Yt7r25gt9pITZVoPxxTMUDlHfdZgyMFfPLKLOtUk3gaiv88FD1Wd4AdI2pNnmYqAMKjqK35Ys3wVuBmgAbUPH6k4ae0ku7/qyRrkGHwRewnAAVZQgZavvBMFJ8cB4Ndb2bXEI2phv40TLtxqffTyNrTfbyrEdKKESwv/h/+pkLpi4bcaUZx0fBeRewm9oUozrsfuGNZcn0rKC8DdMUtqqXBWClU0bAHPTp2GD2c1FNwcwAFOXWyuK3ox0nFYLckD+pQW3ZJyKBKlCnanFXhVDRONHgupmm89QapXZ32FhAzncd/Xcy34ljiSee3UVHJtnq2YJJwQQVBNQKjuhoCv0YqzT/zVWuTbF5QtRm8I7cC5FQANj9TfJqwOxdaoY98K6RRZGzAezoRsF02Q3PwAowC6vYqUzHv85kEfb7BYqzoXGLm4fLxafznJZ+JhJdtsmV5qQiAxjnFQdrdTtL4KNZCXlvtdVYXxLGFN9os7WJoU2kLQB5RTU1rvAlYHfbTZZTTYFKXxRHAizqcSNldKkms7p4RYVAWlUXIGdWTb4kym3qAF4clO+3Njf9F3zxl1b73YfH8pXbnk9fIQnHRowK9vQFRcfD8ikeIUAmZrhe+3cOXZUAt7EeSjZgZFxWFhFll/ZC/Y93aYwMJjAhvbgcS/zZaR0ucmJ663gDu7xgXT3tgQrvb4daDFM10U5HZ/kt2r9m26me5oXdgom65hEkEgfTe03ZFzZnPz7n1fzuGVThCppMAAAyK2pIlNAb40xyGcWNko8VyJ3oghsue2leWHkJ6UpWCdZB13x+QcxkLOr5a8X0zx6N+VcXwNMfGGqlelK0sh1MrRizk4drf0prMW3Q264J3IVJTZ3tx1ULpBi8CnnmrwBeYvfkI5rZJ2W6cOv3bKMqoCpyoNo+gpCZcSVgB8hdzlfMfPl67dL9/IqNo2mBQBaVitoQ4y7KzqHnrBWSxLIyznP3v/VNlzgF7qcPaQiDYhyGAgLfihsZE5nQZLcedY4Apm3dhEOyMUpdzNNyViCP6ykfFn35Ss3JDUS7lfZV47bpxeatoe2dNzU5TGxY9GoVYvhKXmP0WjgOpwv91XPF/x1ty0JRgNGVMwnwxk4PpSpDgGRu/k4jMev9x1vlc3EdAIz+WBCWU23SZoIT3WfyVNpXi7azEQkSB7e1K+fS25uGiNScc4/Nn8BVsUqtwtGhM5nncB7tH0F/k0ngZeGmqEGNJw01v2UYpLQoxwqhCz7/X5it8liVJMxk2kQ3iCGYnwDwfHcl7gLJGiAqrjLBOlp17bPGYIGTlFw+TOCLsGS103TZvSbyptAazsca/TY3o5hZLqX7SBvnfL6s/SU7BVyQY4PpZwHLa6PI2w+uFmJZbBsaruCVPW+FL+zSfvPO560ABJLD6nKsReNBwoxShnNi1XPSmX+C7Ij/ik2nfzubk6IFlomrI0bs386ztTMDifWKNO0ZnC+Oau6ZgwE4PwF8AqKtGFd5yqG0aMFq9rvwPb8DZCXmUHN5Uc40cNG5HWNSAK1yFAkCcG2GPV/7JYEi3t9ZA8NBeh9Dmvsjv8jiaeknPCij14x6BmGOFzyVNqRcS8x517fM39S0cC6bzu8wbnh96fwQ9z/JzCEsqxV9KjuES+xv3rYkX+r3JxcK7p78EACPip85gRBoIU7cZgrwLepMArrpA9YvR8Icag3YILzWcgXVoVRTs6Xc2PsiKasZ98y5CBkmP9Ri3lwZ9acgvPJ1fw729yDzDl/0qEeBOf7ZYZ7dreaPeaGHIymfgKAwmwhnjvAGaqLWhgAAAABJRU5ErkJggg==", - "detail": "high", - }, - }, - ], - }, - ], - "stream": True, - "stream_options": {"include_usage": True}, - "allowed_openai_params": [], - }, - response=SseResponse( - content_type="text/event-stream", - frames=sse_frames( - { - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 5fca57bc9ae9"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "modelVersion": "gemini-3.8-flash", - } - ), - ), - expected=RecountExpected(recount=RecountRates(input_cost_per_token=5e-07, output_cost_per_token=3e-06)), -) GEMINI_GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( name="gemini-gemini-3.8-flash-prompt_blocked", @@ -3543,6 +3273,51 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( ), ) +GEMINI_GEMINI_4_ARGON_INPUT_TEXT: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-4-argon-input_text", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_4_ARGON_STREAM: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-4-argon-stream", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_4_ARGON_CACHE_READ: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-4-argon-cache_read", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_GEMINI_4_ARGON_REASONING: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-4-argon-reasoning", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-4-argon-tool_call", "gemini/gemini-4-argon", + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-3.1-pro-preview-input_text", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-3.1-pro-preview-stream", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-3.1-pro-preview-cache_read", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-3.1-pro-preview-reasoning", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-3.1-pro-preview-tool_call", + "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_INPUT_TEXT, @@ -3570,10 +3345,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_GEMINI_3_8_FLASH_CACHE_READ, - GEMINI_GEMINI_3_8_FLASH_AUDIO_INPUT, - GEMINI_GEMINI_3_8_FLASH_AUDIO_OUTPUT, - GEMINI_GEMINI_3_8_FLASH_IMAGE_INPUT, - GEMINI_GEMINI_3_8_FLASH_VIDEO_INPUT, GEMINI_GEMINI_3_8_FLASH_REASONING, GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, @@ -3582,7 +3353,6 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_8_FLASH_STREAM, GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE, GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL, - GEMINI_GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT, GEMINI_GEMINI_3_8_FLASH_PROMPT_BLOCKED, GEMINI_GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED, GEMINI_GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE, @@ -3605,6 +3375,16 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_NANO_BANANA_2_1_WEB_SEARCH_SINGLE_QUERY, GEMINI_NANO_BANANA_2_1_WEB_SEARCH_MULTIPLE_QUERIES, GEMINI_GEMINI_3_1_PRO_CACHE_CREATION, + GEMINI_GEMINI_4_ARGON_INPUT_TEXT, + GEMINI_GEMINI_4_ARGON_STREAM, + GEMINI_GEMINI_4_ARGON_CACHE_READ, + GEMINI_GEMINI_4_ARGON_REASONING, + GEMINI_GEMINI_4_ARGON_TOOL_CALL, + GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, + GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM, + GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3626,6 +3406,10 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( streamed=GEMINI_GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE, ), StreamParityTestCase(plain=GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, streamed=GEMINI_GEMINI_3_8_FLASH_STREAM_TOOL_CALL), + StreamParityTestCase(plain=GEMINI_GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_GEMINI_4_ARGON_STREAM), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3638,4 +3422,6 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ), ("/v1/responses", "/v1/messages"), ), + *bridge_pairs((GEMINI_GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), + *bridge_pairs((GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), ) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index c3dc0225d0d..f98dda5fb0c 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -2,7 +2,9 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast + +from pydantic import JsonValue from integration.cost_calculation.cost_tracking_case import ( ALWAYS_OPEN_OFF_PEAK_WINDOW, @@ -20,13 +22,91 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames from integration.cost_calculation.cases.anthropic import ( + CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_HAIKU_5_5_CACHE_READ, + CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_INPUT_TEXT, CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_THINKING, + CLAUDE_HAIKU_5_5_TOOL_CALL, ) +def _azure_ai_claude_usage( + input_tokens: int, + output_tokens: int, + inference_geo: str, + *, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ephemeral_5m_input_tokens: int = 0, + ephemeral_1h_input_tokens: int = 0, + thinking_tokens: int = 0, +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation": { + "ephemeral_5m_input_tokens": ephemeral_5m_input_tokens, + "ephemeral_1h_input_tokens": ephemeral_1h_input_tokens, + }, + "output_tokens": output_tokens, + "output_tokens_details": {"thinking_tokens": thinking_tokens}, + "service_tier": "standard", + "inference_geo": inference_geo, + "speed": "standard", + } + + +def _azure_ai_claude_shape_case( + sibling: CostTrackingTestCase, + name: str, + model: str, + usage: dict[str, JsonValue], + expected: ExactExpected, + response_content: list[dict[str, JsonValue]] | None = None, + adaptive_thinking: bool = False, +) -> CostTrackingTestCase: + response: Final = cast(JsonResponse, sibling.response) + request: Final = ( + { + **{ + key: value + for key, value in sibling.request.items() + if key not in {"reasoning_effort", "allowed_openai_params"} + }, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + "allowed_openai_params": ["thinking", "output_config"], + } + if adaptive_thinking + else sibling.request + ) + body: Final = { + **response.body, + "model": model.rsplit("/", 1)[-1], + "usage": usage, + **({"content": response_content} if response_content is not None else {}), + } + return sibling.model_copy( + update={ + "name": name, + "model": model, + "request": request, + "response": response.model_copy(update={"body": body}), + "expected": expected, + } + ) + + # nova-next NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND: Final = CostTrackingTestCase( name="nova-next-transcriptions-per-second", @@ -754,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", @@ -1457,6 +1592,7 @@ AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_INPUT_TEXT.model_ "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, ), @@ -1495,6 +1631,7 @@ AZURE_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, }, @@ -1582,6 +1719,7 @@ AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = CLAUDE_HAIKU_5_5_INPUT_TEXT.model_ "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -1624,6 +1762,7 @@ AZURE_AI_CLAUDE_HAIKU_5_5_STREAM: Final = CLAUDE_HAIKU_5_5_STREAM.model_copy( }, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, }, @@ -1713,6 +1852,157 @@ PERPLEXITY_SONAR_NEXT_OFF_PEAK: Final = PERPLEXITY_SONAR_NEXT_NO_SEARCH.model_co } ) + +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_READ, + "azure_ai/claude-fable-5-1-cache_read", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=14, output_tokens=6, inference_geo="not_available", cache_read_input_tokens=27003), + ExactExpected( + spend=0.012224275, + input_cost=0.011714275, + output_cost=0.00051, + prompt_tokens=27017, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + "azure_ai/claude-fable-5-1-cache_write_5m", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=5, + inference_geo="not_available", + cache_creation_input_tokens=27011, + ephemeral_5m_input_tokens=27011, + ), + ExactExpected( + spend=0.57464675, + input_cost=0.57422175, + output_cost=0.000425, + prompt_tokens=27025, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + "azure_ai/claude-fable-5-1-cache_write_1h", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=5, + inference_geo="not_available", + cache_creation_input_tokens=27011, + ephemeral_1h_input_tokens=27011, + ), + ExactExpected( + spend=0.919037, + input_cost=0.918612, + output_cost=0.000425, + prompt_tokens=27025, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_FABLE_5_1_THINKING: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_THINKING, + "azure_ai/claude-fable-5-1-thinking", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=18, output_tokens=99, inference_geo="not_available", thinking_tokens=33), + ExactExpected( + spend=0.008721, input_cost=0.000306, output_cost=0.008415, prompt_tokens=18, completion_tokens=99 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _azure_ai_claude_shape_case( + CLAUDE_FABLE_5_1_TOOL_CALL, + "azure_ai/claude-fable-5-1-tool_call", + "azure_ai/claude-fable-5-1", + _azure_ai_claude_usage(input_tokens=385, output_tokens=50, inference_geo="not_available"), + ExactExpected( + spend=0.010795, input_cost=0.006545, output_cost=0.00425, prompt_tokens=385, completion_tokens=50 + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_READ, + "azure_ai/claude-haiku-5-5-cache_read", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=14, output_tokens=5, inference_geo="global", cache_read_input_tokens=27003), + ExactExpected( + spend=0.000520467, + input_cost=0.000515717, + output_cost=0.00000475, + prompt_tokens=27017, + completion_tokens=5, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + "azure_ai/claude-haiku-5-5-cache_write_5m", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=6, + inference_geo="global", + cache_creation_input_tokens=27011, + ephemeral_5m_input_tokens=27011, + ), + ExactExpected( + spend=0.0064234725, + input_cost=0.0064177725, + output_cost=0.0000057, + prompt_tokens=27025, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + "azure_ai/claude-haiku-5-5-cache_write_1h", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage( + input_tokens=14, + output_tokens=6, + inference_geo="global", + cache_creation_input_tokens=27011, + ephemeral_1h_input_tokens=27011, + ), + ExactExpected( + spend=0.01027254, + input_cost=0.01026684, + output_cost=0.0000057, + prompt_tokens=27025, + completion_tokens=6, + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_THINKING: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_THINKING, + "azure_ai/claude-haiku-5-5-thinking", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=18, output_tokens=95, inference_geo="global", thinking_tokens=54), + ExactExpected( + spend=0.00009367, input_cost=0.00000342, output_cost=0.00009025, prompt_tokens=18, completion_tokens=95 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _azure_ai_claude_shape_case( + CLAUDE_HAIKU_5_5_TOOL_CALL, + "azure_ai/claude-haiku-5-5-tool_call", + "azure_ai/claude-haiku-5-5", + _azure_ai_claude_usage(input_tokens=511, output_tokens=34, inference_geo="global"), + ExactExpected( + spend=0.00012939, input_cost=0.00009709, output_cost=0.0000323, prompt_tokens=511, completion_tokens=34 + ), +) + + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( PERPLEXITY_SONAR_NEXT_OFF_PEAK, NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND, @@ -1743,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, @@ -1764,8 +2056,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( XAI_GROK_4_7_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_STREAM, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_READ, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + AZURE_AI_CLAUDE_FABLE_5_1_THINKING, + AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL, AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + AZURE_AI_CLAUDE_HAIKU_5_5_THINKING, + AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index cfb7fda4d71..85ac05dd69d 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -26,6 +26,14 @@ from integration.cost_calculation.stream_parity.case import StreamParityTestCase from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs +def _response_for_model(response: JsonResponse, model: str) -> JsonResponse: + return response.model_copy(update={"body": {**response.body, "model": model}}) + + +def _stream_for_model(response: SseResponse, model: str) -> SseResponse: + return response.model_copy(update={"frames": tuple(frame.replace("gpt-6-luna", model) for frame in response.frames)}) + + # gpt-5.3-codex GPT_5_3_CODEX_INPUT_TEXT: Final = CostTrackingTestCase( name="gpt-5.3-codex-input_text", @@ -11183,6 +11191,171 @@ GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( ExactExpected(spend=0.2148, input_cost=0.192, output_cost=0.0228, prompt_tokens=12928, completion_tokens=380), ) +GPT_5_6_TERRA_INPUT_TEXT: Final = GPT_6_LUNA_INPUT_TEXT.model_copy( + update={ + "name": "gpt-5.6-terra-input_text", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_INPUT_TEXT.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_STREAM: Final = GPT_6_LUNA_STREAM.model_copy( + update={ + "name": "gpt-5.6-terra-stream", + "model": "gpt-5.6-terra", + "response": _stream_for_model(cast(SseResponse, GPT_6_LUNA_STREAM.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_SERVICE_TIER_FLEX: Final = GPT_6_LUNA_SERVICE_TIER_FLEX.model_copy( + update={ + "name": "gpt-5.6-terra-service_tier_flex", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_SERVICE_TIER_FLEX.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0056056, + input_cost=0.002392, + output_cost=0.0032136, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_SERVICE_TIER_PRIORITY: Final = GPT_6_LUNA_SERVICE_TIER_PRIORITY.model_copy( + update={ + "name": "gpt-5.6-terra-service_tier_priority", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_SERVICE_TIER_PRIORITY.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0224224, + input_cost=0.009568, + output_cost=0.0128544, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_RESPONSES_NATIVE_JSON: Final = GPT_6_LUNA_RESPONSES_NATIVE_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-responses_native_json", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_RESPONSES_NATIVE_JSON.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON: Final = GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-responses_deployment_json", + "model": "gpt-5.6-terra", + "deployment": Deployment(model="openai/responses/gpt-5.6-terra"), + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_CHAT_BRIDGED_JSON: Final = GPT_6_LUNA_CHAT_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-chat_bridged_json", + "model": "gpt-5.6-terra", + "deployment": Deployment(model="openai/responses/gpt-5.6-terra"), + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_CHAT_BRIDGED_JSON.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) +GPT_5_6_TERRA_CACHE_READ: Final = GPT_6_LUNA_CACHE_READ.model_copy( + update={ + "name": "gpt-5.6-terra-cache_read", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_CACHE_READ.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.01078688, + input_cost=0.00485888, + output_cost=0.005928, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +GPT_5_6_TERRA_CACHE_WRITE_5M: Final = GPT_6_LUNA_CACHE_WRITE_5M.model_copy( + update={ + "name": "gpt-5.6-terra-cache_write_5m", + "model": "gpt-5.6-terra", + "response": _response_for_model( + cast(JsonResponse, GPT_6_LUNA_CACHE_WRITE_5M.response), "gpt-5.6-terra" + ), + "expected": ExactExpected( + spend=0.047528, + input_cost=0.0416, + output_cost=0.005928, + prompt_tokens=12928, + completion_tokens=380, + ), + } +) +GPT_5_6_TERRA_REASONING: Final = GPT_6_LUNA_REASONING.model_copy( + update={ + "name": "gpt-5.6-terra-reasoning", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_REASONING.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.066248, + input_cost=0.003224, + output_cost=0.063024, + prompt_tokens=1240, + completion_tokens=4040, + ), + } +) +GPT_5_6_TERRA_TOOL_CALL: Final = GPT_6_LUNA_TOOL_CALL.model_copy( + update={ + "name": "gpt-5.6-terra-tool_call", + "model": "gpt-5.6-terra", + "response": _response_for_model(cast(JsonResponse, GPT_6_LUNA_TOOL_CALL.response), "gpt-5.6-terra"), + "expected": ExactExpected( + spend=0.0112112, + input_cost=0.004784, + output_cost=0.0064272, + prompt_tokens=1840, + completion_tokens=412, + ), + } +) + GPT_5_6_SLASH_ALIAS: Final = GPT_5_6_INPUT_TEXT.model_copy( update={"name": "gpt-5.6-slash_alias", "alias": "openai/custom/gpt-5.6-$MARKER"} @@ -11485,6 +11658,17 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_6_1_SOL_CACHE_WRITE_5M, GPT_6_LUNA_CACHE_WRITE_5M, GPT_6_ASTRA_CACHE_WRITE_5M, + GPT_5_6_TERRA_INPUT_TEXT, + GPT_5_6_TERRA_STREAM, + GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, + GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, + GPT_5_6_TERRA_CHAT_BRIDGED_JSON, + GPT_5_6_TERRA_CACHE_READ, + GPT_5_6_TERRA_CACHE_WRITE_5M, + GPT_5_6_TERRA_REASONING, + GPT_5_6_TERRA_SERVICE_TIER_FLEX, + GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, + GPT_5_6_TERRA_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -11513,6 +11697,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_6_LUNA_INPUT_TEXT, streamed=GPT_6_LUNA_STREAM), StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), + StreamParityTestCase(plain=GPT_5_6_TERRA_INPUT_TEXT, streamed=GPT_5_6_TERRA_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -11545,4 +11730,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_LUNA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_ASTRA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_1_SOL_CHAT_BRIDGED_JSON), + BridgeCostParityTestCase( + native=GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_5_6_TERRA_CHAT_BRIDGED_JSON + ), ) diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 966a24bcba2..761078ef593 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2,15 +2,32 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" -from typing import Final +from typing import Final, cast + +from pydantic import JsonValue from integration.cost_calculation.cases.anthropic import ( + CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_OPUS_5_5_CACHE_READ, + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_INPUT_TEXT, CLAUDE_OPUS_5_5_STREAM, + CLAUDE_OPUS_5_5_THINKING, + CLAUDE_OPUS_5_5_TOOL_CALL, + CLAUDE_SONNET_5_5_CACHE_READ, + CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_INPUT_TEXT, CLAUDE_SONNET_5_5_STREAM, + CLAUDE_SONNET_5_5_THINKING, + CLAUDE_SONNET_5_5_TOOL_CALL, ) from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( @@ -24,6 +41,99 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: + match response: + case JsonResponse(): + return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + case SseResponse(): + return response.model_copy( + update={ + "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + } + ) + + +def _gemini_case( + sibling: CostTrackingTestCase, name: str, model: str, expected: ExactExpected +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "model": model, + "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), "gemini-3.1-pro-preview"), + "expected": expected, + } + ) + + +def _vertex_claude_usage( + input_tokens: int, + output_tokens: int, + *, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ephemeral_5m_input_tokens: int = 0, + ephemeral_1h_input_tokens: int = 0, + thinking_tokens: int = 0, +) -> dict[str, JsonValue]: + return { + "input_tokens": input_tokens, + "cache_creation_input_tokens": cache_creation_input_tokens, + "cache_read_input_tokens": cache_read_input_tokens, + "cache_creation": { + "ephemeral_5m_input_tokens": ephemeral_5m_input_tokens, + "ephemeral_1h_input_tokens": ephemeral_1h_input_tokens, + }, + "output_tokens": output_tokens, + "output_tokens_details": {"thinking_tokens": thinking_tokens}, + } + + +def _vertex_claude_shape_case( + sibling: CostTrackingTestCase, + name: str, + model: str, + usage: dict[str, JsonValue], + expected: ExactExpected, + response_content: list[dict[str, JsonValue]] | None = None, + adaptive_thinking: bool = False, +) -> CostTrackingTestCase: + response: Final = cast(JsonResponse, sibling.response) + request: Final = ( + { + **{ + key: value + for key, value in sibling.request.items() + if key not in {"reasoning_effort", "allowed_openai_params"} + }, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium"}, + "allowed_openai_params": ["thinking", "output_config"], + } + if adaptive_thinking + else sibling.request + ) + body: Final = { + key: value + for key, value in { + **response.body, + "model": model.rsplit("/", 1)[-1], + "usage": usage, + **({"content": response_content} if response_content is not None else {}), + }.items() + if key != "diagnostics" + } + return sibling.model_copy( + update={ + "name": name, + "model": model, + "request": request, + "response": response.model_copy(update={"body": body}), + "expected": expected, + } + ) + + def _vertex_claude_stream_response(model: str, stop_reason: str, text_deltas: tuple[str, ...]) -> SseResponse: text_frames: Final = ( ( @@ -138,7 +248,9 @@ GEMINI_3_1_PRO_INPUT_TEXT: Final = CostTrackingTestCase( "promptTokenCount": 1840, "candidatesTokenCount": 412, "totalTokenCount": 2252, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], }, "modelVersion": "gemini-3.1-pro", }, @@ -192,7 +304,9 @@ GEMINI_3_1_PRO_CACHE_READ: Final = CostTrackingTestCase( "candidatesTokenCount": 380, "totalTokenCount": 13308, "cachedContentTokenCount": 12288, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], }, "modelVersion": "gemini-3.1-pro", }, @@ -825,7 +939,9 @@ GEMINI_3_1_PRO_STREAM: Final = CostTrackingTestCase( "promptTokenCount": 1840, "candidatesTokenCount": 412, "totalTokenCount": 2252, + "trafficType": "ON_DEMAND", "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], }, "modelVersion": "gemini-3.1-pro", }, @@ -1557,9 +1673,9 @@ GEMINI_3_1_PRO_PASSTHROUGH_STREAM_GENERATE_CONTENT_PRICED_VIA_VERTEX_KEY: Final # gemini-3.8-flash GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-input_text", + name="vertex_ai/gemini-3.8-flash-input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1610,9 +1726,9 @@ GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( - name="gemini-3.8-flash-cache_read", + name="vertex_ai/gemini-3.8-flash-cache_read", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1663,196 +1779,13 @@ GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( ), ) -GEMINI_3_8_FLASH_AUDIO_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-audio_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "3e33892c4f9f summarize the attached material in one line and name the city weather", - }, - { - "type": "input_audio", - "input_audio": { - "data": "UklGRmQGAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YUAGAAAAAA8I4A85F+EdpCNYKNkrCy7eLkwuWSwTKZUk/x59GEERgQl4AWb5hPER6kbjVd1t2LLUQdIu0X/RM9M81oTa6t9E5mPtD/UP/SQFEw2fFI0bqSHEJrgqZi27LqouNi1pKlkmJSH0GvUTXgxpBFP8WPS27KflYt8U2ujV/dJp0TjRbNL71NPY1d3b47nqOPIh+jUCOQrwER0ZjB8JJW0plCxoLtou5i2VK/cnKSNOHZUWLw9VB0T/OPdv7yToj+Hi20nX5tPT0SDR09Hm00nX4tuP4STob+8490T/VQcvD5UWTh0pI/cnlSvmLdouaC6ULG0pCSWMHx0Z8BE5CjUCIfo48rnq2+PV3dPY+9Rs0jjRadH90ujVFNpi36fltuxY9FP8aQReDPUT9BolIVkmaSo2Laouuy5mLbgqxCapIY0bnxQTDSQFD/0P9WPtRObq34TaPNYz03/RLtFB0rLUbdhV3UbjEeqE8Wb5eAGBCUERfRj/HpUkEylZLEwu3i4LLtkrWCikI+EdORfgDw8IAADx9yDwx+gf4lzcqNcn1PXRItG00afT7dZr2wHhg+e/7n/2iP6aBnwO7xW6HKsikydOK78t0i6BLs0sxCl8JRYgvBmdEvEK8QLc+u3yYetz5FfePNlI1ZrSRdFW0crSl9Wn2dveDOUL7KLzl/utA6gLShNZGp4g7CUYKgMtly7ILpQtBSstJysiJRxHFcgN3wXL/cf1EO7j5nTg99qT1mzTmNEm0RrSa9QJ2NfcsuJr6dHwq/i8AMgIkRDcF3EeHiS3KBosLS7gLi0uGiy3KB4kcR7cF5EQyAi8AKv40fBr6bLi19wJ2GvUGtIm0ZjRbNOT1vfadODj5hDux/XL/d8FyA1HFSUcKyItJwUrlC3ILpcuAy0YKuwlniBZGkoTqAutA5f7ovML7Azl296n2ZfVytJW0UXRmtJI1TzZV95z5GHr7fLc+vEC8QqdErwZFiB8JcQpzSyBLtIuvy1OK5MnqyK6HO8VfA6aBoj+f/a/7oPnAeFr2+3Wp9O00SLR9dEn1KjXXNwf4sfoIPDx9wAADwjgDzkX4R2kI1go2SsLLt4uTC5ZLBMplST/Hn0YQRGBCXgBZvmE8RHqRuNV3W3YstRB0i7Rf9Ez0zzWhNrq30TmY+0P9Q/9JAUTDZ8UjRupIcQmuCpmLbsuqi42LWkqWSYlIfQa9RNeDGkEU/xY9Lbsp+Vi3xTa6NX90mnRONFs0vvU09jV3dvjueo48iH6NQI5CvARHRmMHwklbSmULGgu2i7mLZUr9ycpI04dlRYvD1UHRP8492/vJOiP4eLbSdfm09PRINHT0ebTSdfi24/hJOhv7zj3RP9VBy8PlRZOHSkj9yeVK+Yt2i5oLpQsbSkJJYwfHRnwETkKNQIh+jjyuerb49Xd09j71GzSONFp0f3S6NUU2mLfp+W27Fj0U/xpBF4M9RP0GiUhWSZpKjYtqi67LmYtuCrEJqkhjRufFBMNJAUP/Q/1Y+1E5urfhNo81jPTf9Eu0UHSstRt2FXdRuMR6oTxZvl4AYEJQRF9GP8elSQTKVksTC7eLgsu2StYKKQj4R05F+APDwgAAPH3IPDH6B/iXNyo1yfU9dEi0bTRp9Pt1mvbAeGD57/uf/aI/poGfA7vFbocqyKTJ04rvy3SLoEuzSzEKXwlFiC8GZ0S8QrxAtz67fJh63PkV9482UjVmtJF0VbRytKX1afZ294M5QvsovOX+60DqAtKE1kaniDsJRgqAy2XLsgulC0FKy0nKyIlHEcVyA3fBcv9x/UQ7uPmdOD32pPWbNOY0SbRGtJr1AnY19yy4mvp0fCr+LwAyAiRENwXcR4eJLcoGiwtLuAuLS4aLLcoHiRxHtwXkRDICLwAq/jR8GvpsuLX3AnYa9Qa0ibRmNFs05PW99p04OPmEO7H9cv93wXIDUcVJRwrIi0nBSuULcguly4DLRgq7CWeIFkaShOoC60Dl/ui8wvsDOXb3qfZl9XK0lbRRdGa0kjVPNlX3nPkYevt8tz68QLxCp0SvBkWIHwlxCnNLIEu0i6/LU4rkyerIroc7xV8DpoGiP5/9r/ug+cB4Wvb7dan07TRItH10SfUqNdc3B/ix+gg8PH3", - "format": "wav", - }, - }, - ], - }, - ], - "stream": False, - "modalities": ["text"], - "allowed_openai_params": ["modalities"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 3e33892c4f9f"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1546, - "candidatesTokenCount": 210, - "totalTokenCount": 1756, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 96}, - {"modality": "AUDIO", "tokenCount": 1450}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.00221312, input_cost=0.00155792, output_cost=0.0006552, prompt_tokens=1546, completion_tokens=210 - ), -) -GEMINI_3_8_FLASH_AUDIO_OUTPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-audio_output", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "30bf1c0de6fe summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "modalities": ["text", "audio"], - "audio": {"voice": "alloy", "format": "pcm16"}, - "allowed_openai_params": ["modalities", "audio"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 30bf1c0de6fe"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 220, - "candidatesTokenCount": 1300, - "totalTokenCount": 1520, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 220}], - "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 180}, - {"modality": "AUDIO", "tokenCount": 1120}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0076648, input_cost=0.0001144, output_cost=0.0075504, prompt_tokens=220, completion_tokens=1300 - ), -) -GEMINI_3_8_FLASH_VIDEO_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-video_input", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "5f5da5957185 summarize the attached material in one line and name the city weather", - }, - { - "type": "file", - "file": { - "file_data": "data:video/mp4;base64,AAAAGGZ0eXBpc29tAAACAGlzb21pc282AAACCG1kYXQNFBsiKTA3PkVMU1phaG92fYSLkpmgp661vMPK0djf5u30+wIJEBceJSwzOkFIT1ZdZGtyeYCHjpWco6qxuL/GzdTb4unw9/4FDBMaISgvNj1ES1JZYGdudXyDipGYn6attLvCydDX3uXs8/oBCA8WHSQrMjlAR05VXGNqcXh/ho2Um6KpsLe+xczT2uHo7/b9BAsSGSAnLjU8Q0pRWF9mbXR7gomQl56lrLO6wcjP1t3k6/L5AAcOFRwjKjE4P0ZNVFtiaXB3foWMk5qhqK+2vcTL0tng5+71/AMKERgfJi00O0JJUFdeZWxzeoGIj5adpKuyucDHztXc4+rx+P8GDRQbIikwNz5FTFNaYWhvdn2Ei5KZoKeutbzDytHY3+bt9PsCCRAXHiUsMzpBSE9WXWRrcnmAh46VnKOqsbi/xs3U2+Lp8Pf+BQwTGiEoLzY9REtSWWBnbnV8g4qRmJ+mrbS7wsnQ197l7PP6AQgPFh0kKzI5QEdOVVxjanF4f4aNlJuiqbC3vsXM09rh6O/2/QQLEhkgJy41PENKUVhfZm10e4KJkJeepayzusHIz9bd5Ovy+QAHDhUcIyoxOD9GTVRbYmlwd36FjJOaoaivtr3Ey9LZ4Ofu9fwDChEYHyYtNDtCSVBXXmVsc3qBiI+WnaSrsrnAx87V3OPq8fj/Bg==", - "format": "mp4", - }, - }, - ], - }, - ], - "stream": False, - "allowed_openai_params": [], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer 5f5da5957185"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 8060, - "candidatesTokenCount": 300, - "totalTokenCount": 8360, - "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 140}, - {"modality": "VIDEO", "tokenCount": 7920}, - ], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0059192, input_cost=0.0049832, output_cost=0.000936, prompt_tokens=8060, completion_tokens=300 - ), -) GEMINI_3_8_FLASH_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( - name="gemini-3.8-flash-service_tier_flex", + name="vertex_ai/gemini-3.8-flash-service_tier_flex", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1904,9 +1837,9 @@ GEMINI_3_8_FLASH_SERVICE_TIER_FLEX: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( - name="gemini-3.8-flash-service_tier_priority", + name="vertex_ai/gemini-3.8-flash-service_tier_priority", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -1958,9 +1891,9 @@ GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-web_search_per_prompt", + name="vertex_ai/gemini-3.8-flash-web_search_per_prompt", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2008,14 +1941,14 @@ GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT: Final = CostTrackingTestCase( }, ), expected=ExactExpected( - spend=0.03724224, input_cost=0.0009568, output_cost=0.00128544, prompt_tokens=1840, completion_tokens=412 + spend=0.10724224, input_cost=0.0009568, output_cost=0.00128544, prompt_tokens=1840, completion_tokens=412 ), ) GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING: Final = CostTrackingTestCase( - name="gemini-3.8-flash-google_maps_grounding", + name="vertex_ai/gemini-3.8-flash-google_maps_grounding", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2071,65 +2004,10 @@ GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING: Final = CostTrackingTestCase( ), ) -GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-fallback_reasoning_at_output_rate", - covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", - request={ - "model": "$MODEL", - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", - } - ], - }, - { - "role": "user", - "content": [ - { - "type": "text", - "text": "f92946792f44 summarize the attached material in one line and name the city weather", - } - ], - }, - ], - "stream": False, - "reasoning_effort": "medium", - "allowed_openai_params": ["reasoning_effort"], - }, - response=JsonResponse( - content_type="application/json", - body={ - "candidates": [ - { - "content": {"parts": [{"text": "scripted answer f92946792f44"}], "role": "model"}, - "finishReason": "STOP", - "index": 0, - } - ], - "usageMetadata": { - "promptTokenCount": 1240, - "candidatesTokenCount": 560, - "thoughtsTokenCount": 3480, - "totalTokenCount": 5280, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], - }, - "modelVersion": "gemini-3.8-flash", - }, - ), - expected=ExactExpected( - spend=0.0132496, input_cost=0.0006448, output_cost=0.0126048, prompt_tokens=1240, completion_tokens=4040 - ), -) - GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-fallback_image_tokens_at_input_rate", + name="vertex_ai/gemini-3.8-flash-fallback_image_tokens_at_input_rate", covers="quota_management.spend_tracking.cost_matrix.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2190,9 +2068,9 @@ GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE: Final = CostTrackingTestCa ) GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream", + name="vertex_ai/gemini-3.8-flash-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2250,9 +2128,9 @@ GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage", + name="vertex_ai/gemini-3.8-flash-stream_no_usage", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2302,9 +2180,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage_tool_call", + name="vertex_ai/gemini-3.8-flash-stream_no_usage_tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2386,9 +2264,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_no_usage_image_input", + name="vertex_ai/gemini-3.8-flash-stream_no_usage_image_input", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2445,9 +2323,9 @@ GEMINI_3_8_FLASH_STREAM_NO_USAGE_IMAGE_INPUT: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( - name="gemini-3.8-flash-prompt_blocked", + name="vertex_ai/gemini-3.8-flash-prompt_blocked", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2495,9 +2373,9 @@ GEMINI_3_8_FLASH_PROMPT_BLOCKED: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_prompt_blocked", + name="vertex_ai/gemini-3.8-flash-stream_prompt_blocked", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2552,9 +2430,9 @@ GEMINI_3_8_FLASH_STREAM_PROMPT_BLOCKED: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-response_model_override", + name="vertex_ai/gemini-3.8-flash-response_model_override", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2605,9 +2483,9 @@ GEMINI_3_8_FLASH_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_response_model_override", + name="vertex_ai/gemini-3.8-flash-stream_response_model_override", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2665,9 +2543,9 @@ GEMINI_3_8_FLASH_STREAM_RESPONSE_MODEL_OVERRIDE: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-tool_call", + name="vertex_ai/gemini-3.8-flash-tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2752,9 +2630,9 @@ GEMINI_3_8_FLASH_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_TOOL_CALL: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_tool_call", + name="vertex_ai/gemini-3.8-flash-stream_tool_call", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2846,9 +2724,9 @@ GEMINI_3_8_FLASH_STREAM_TOOL_CALL: Final = CostTrackingTestCase( ) GEMINI_3_8_FLASH_STREAM_FULL_USAGE: Final = CostTrackingTestCase( - name="gemini-3.8-flash-stream_full_usage", + name="vertex_ai/gemini-3.8-flash-stream_full_usage", covers="quota_management.spend_tracking.scripted_wire.logs_cost", - model="gemini-3.8-flash", + model="vertex_ai/gemini-3.8-flash", request={ "model": "$MODEL", "messages": [ @@ -2910,7 +2788,7 @@ GEMINI_3_8_FLASH_STREAM_FULL_USAGE: Final = CostTrackingTestCase( ), ), expected=ExactExpected( - spend=0.007460128, input_cost=0.001619488, output_cost=0.00584064, prompt_tokens=8314, completion_tokens=1592 + spend=0.005775328, input_cost=0.001276288, output_cost=0.00449904, prompt_tokens=8314, completion_tokens=1592 ), ) @@ -3081,6 +2959,347 @@ VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( } ) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_READ, + "vertex_ai/claude-sonnet-5-5-cache_read", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.00474011, input_cost=0.00463811, output_cost=0.000102, prompt_tokens=27017, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-sonnet-5-5-cache_write_5m", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.11494635, input_cost=0.11484435, output_cost=0.000102, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-sonnet-5-5-cache_write_1h", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.1838244, input_cost=0.1837224, output_cost=0.000102, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_SONNET_5_5_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_THINKING, + "vertex_ai/claude-sonnet-5-5-thinking", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=144, thinking_tokens=37), + ExactExpected( + spend=0.0025092, input_cost=0.0000612, output_cost=0.002448, prompt_tokens=18, completion_tokens=144 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_SONNET_5_5_TOOL_CALL, + "vertex_ai/claude-sonnet-5-5-tool_call", + "vertex_ai/claude-sonnet-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.002159, input_cost=0.001309, output_cost=0.00085, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_READ, + "vertex_ai/claude-opus-5-5-cache_read", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.01003788, input_cost=0.00982188, output_cost=0.000216, prompt_tokens=27017, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-opus-5-5-cache_write_5m", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.2434158, input_cost=0.2431998, output_cost=0.000216, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-opus-5-5-cache_write_1h", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.3892752, input_cost=0.3890592, output_cost=0.000216, prompt_tokens=27025, completion_tokens=6 + ), +) +VERTEX_AI_CLAUDE_OPUS_5_5_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_THINKING, + "vertex_ai/claude-opus-5-5-thinking", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=162, thinking_tokens=66), + ExactExpected( + spend=0.0059616, input_cost=0.0001296, output_cost=0.005832, prompt_tokens=18, completion_tokens=162 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_OPUS_5_5_TOOL_CALL, + "vertex_ai/claude-opus-5-5-tool_call", + "vertex_ai/claude-opus-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.004572, input_cost=0.002772, output_cost=0.0018, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_READ, + "vertex_ai/claude-fable-5-1-cache_read", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=14, output_tokens=5, cache_read_input_tokens=27003), + ExactExpected( + spend=0.013567425, input_cost=0.013092425, output_cost=0.000475, prompt_tokens=27017, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + "vertex_ai/claude-fable-5-1-cache_write_5m", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.64225225, input_cost=0.64177725, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + "vertex_ai/claude-fable-5-1-cache_write_1h", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=1.027159, input_cost=1.026684, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 + ), +) +VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_THINKING, + "vertex_ai/claude-fable-5-1-thinking", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=18, output_tokens=98, thinking_tokens=33), + ExactExpected( + spend=0.009652, input_cost=0.000342, output_cost=0.00931, prompt_tokens=18, completion_tokens=98 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) +VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_1_TOOL_CALL, + "vertex_ai/claude-fable-5-1-tool_call", + "vertex_ai/claude-fable-5-1", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.012065, input_cost=0.007315, output_cost=0.00475, prompt_tokens=385, completion_tokens=50 + ), +) + + +GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( + GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( + GEMINI_3_8_FLASH_STREAM, "vertex_ai/gemini-3.1-pro-preview-stream", "vertex_ai/gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) +GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( + GEMINI_3_8_FLASH_CACHE_READ, "vertex_ai/gemini-3.1-pro-preview-cache_read", "vertex_ai/gemini-3.1-pro-preview", + ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) +) +GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( + GEMINI_3_1_PRO_REASONING, "vertex_ai/gemini-3.1-pro-preview-reasoning", "vertex_ai/gemini-3.1-pro-preview", + ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) +) +GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( + GEMINI_3_8_FLASH_TOOL_CALL, "vertex_ai/gemini-3.1-pro-preview-tool_call", "vertex_ai/gemini-3.1-pro-preview", + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) +) + +GEMINI_4_ARGON_INPUT_TEXT: Final = GEMINI_3_8_FLASH_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-input_text", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": {"parts": [{"text": "Hello."}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 43, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.0001608, input_cost=0.0000888, output_cost=0.000072, prompt_tokens=37, completion_tokens=6 + ), + } +) +GEMINI_4_ARGON_STREAM: Final = GEMINI_3_8_FLASH_STREAM.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-stream", + "model": "vertex_ai/gemini-4-argon", + "response": SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hello."}]}}], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-4-argon", + }, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 43, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + ), + "expected": ExactExpected( + spend=0.0001608, input_cost=0.0000888, output_cost=0.000072, prompt_tokens=37, completion_tokens=6 + ), + } +) +GEMINI_4_ARGON_CACHE_READ: Final = GEMINI_3_8_FLASH_CACHE_READ.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-cache_read", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + {"content": {"parts": [{"text": "Understood."}], "role": "model"}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 13532, + "candidatesTokenCount": 21, + "totalTokenCount": 13553, + "cachedContentTokenCount": 11478, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 13532}], + "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 11478}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 21}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.00655896, input_cost=0.00630696, output_cost=0.000252, prompt_tokens=13532, completion_tokens=21 + ), + } +) +GEMINI_4_ARGON_REASONING: Final = GEMINI_3_1_PRO_REASONING.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-reasoning", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": { + "parts": [{"text": "Hello.", "thought": True}, {"text": "Hello."}], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 37, + "candidatesTokenCount": 6, + "totalTokenCount": 162, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 37}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 6}], + "thoughtsTokenCount": 119, + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.0015888, input_cost=0.0000888, output_cost=0.0015, prompt_tokens=37, completion_tokens=125 + ), + } +) +GEMINI_4_ARGON_TOOL_CALL: Final = GEMINI_3_8_FLASH_TOOL_CALL.model_copy( + update={ + "name": "vertex_ai/gemini-4-argon-tool_call", + "model": "vertex_ai/gemini-4-argon", + "response": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": { + "parts": [{"functionCall": {"name": "report_weather", "args": {"city": "Paris"}}}], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 90, + "candidatesTokenCount": 22, + "totalTokenCount": 112, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 90}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 22}], + }, + "modelVersion": "gemini-4-argon", + }, + ), + "expected": ExactExpected( + spend=0.00048, input_cost=0.000216, output_cost=0.000264, prompt_tokens=90, completion_tokens=22 + ), + } +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, @@ -3108,14 +3327,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, - GEMINI_3_8_FLASH_AUDIO_INPUT, - GEMINI_3_8_FLASH_AUDIO_OUTPUT, - GEMINI_3_8_FLASH_VIDEO_INPUT, GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, GEMINI_3_8_FLASH_WEB_SEARCH_PER_PROMPT, GEMINI_3_8_FLASH_GOOGLE_MAPS_GROUNDING, - GEMINI_3_8_FLASH_FALLBACK_REASONING_AT_OUTPUT_RATE, GEMINI_3_8_FLASH_FALLBACK_IMAGE_TOKENS_AT_INPUT_RATE, GEMINI_3_8_FLASH_STREAM, GEMINI_3_8_FLASH_STREAM_NO_USAGE, @@ -3137,6 +3352,31 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_STREAM, VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_STREAM, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_SONNET_5_5_THINKING, + VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_OPUS_5_5_THINKING, + VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, + VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, + GEMINI_3_1_PRO_PREVIEW_STREAM, + GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_3_1_PRO_PREVIEW_REASONING, + GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, + GEMINI_4_ARGON_INPUT_TEXT, + GEMINI_4_ARGON_STREAM, + GEMINI_4_ARGON_CACHE_READ, + GEMINI_4_ARGON_REASONING, + GEMINI_4_ARGON_TOOL_CALL, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3157,9 +3397,15 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), + StreamParityTestCase(plain=GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_4_ARGON_STREAM), ) -BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = bridge_pairs( - (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), - ("/v1/responses", "/v1/messages"), +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + *bridge_pairs( + (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), + ("/v1/responses", "/v1/messages"), + ), + *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), + *bridge_pairs((GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), ) diff --git a/tests/integration/cost_calculation/catalog.py b/tests/integration/cost_calculation/catalog.py index 66eb7497bd7..568a89ad967 100644 --- a/tests/integration/cost_calculation/catalog.py +++ b/tests/integration/cost_calculation/catalog.py @@ -60,6 +60,7 @@ REALTIME_CASES: Final[tuple[RealtimeCostCase, ...]] = batch_realtime.REALTIME_CA BRIDGE_COST_PARITY_CASES: Final[tuple[BridgeCostParityTestCase, ...]] = ( *openai.BRIDGE_COST_PARITY, *bedrock_converse.BRIDGE_COST_PARITY, + *bedrock_mantle.BRIDGE_COST_PARITY, *azure.BRIDGE_COST_PARITY, *anthropic.BRIDGE_COST_PARITY, *gemini.BRIDGE_COST_PARITY, diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 78d025855b5..2ef42249a7e 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -32,6 +32,41 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "input_cost_per_token": 3.2e-06, "output_cost_per_token": 1.6e-05, }, + "gpt-5.6-terra": { + "cache_creation_input_token_cost": 3.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 3.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.3e-05, + "cache_creation_input_token_cost_flex": 1.625e-06, + "cache_creation_input_token_cost_priority": 6.5e-06, + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.2e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.6e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.04e-06, + "cache_read_input_token_cost_flex": 1.3e-07, + "cache_read_input_token_cost_priority": 5.2e-07, + "input_cost_per_token": 2.6e-06, + "input_cost_per_token_above_272k_tokens": 5.2e-06, + "input_cost_per_token_above_272k_tokens_flex": 2.6e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.04e-05, + "input_cost_per_token_flex": 1.3e-06, + "input_cost_per_token_priority": 5.2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.56e-05, + "output_cost_per_token_above_272k_tokens": 2.34e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.17e-05, + "output_cost_per_token_above_272k_tokens_priority": 4.68e-05, + "output_cost_per_token_flex": 7.8e-06, + "output_cost_per_token_priority": 3.12e-05, + "supports_function_calling": True, + "supports_none_reasoning_effort": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, "gpt-6-luna": { "output_cost_per_token_priority": 1e-06, "input_cost_per_token_priority": 2e-07, @@ -407,21 +442,64 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, "web_search_billing_unit": "per_query", }, - "gemini/gemini-3.8-flash": { - "cache_read_input_token_cost": 5e-08, - "google_maps_grounding_cost_per_query": 0.025, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 5.5e-07, - "input_cost_per_token": 5e-07, - "input_cost_per_token_flex": 2.5e-07, - "input_cost_per_token_priority": 6.25e-07, - "input_cost_per_video_token": 6e-07, + "gemini/gemini-4-argon": { + "cache_read_input_token_cost": 1.1e-07, + "cache_read_input_token_cost_flex": 5.5e-08, + "cache_read_input_token_cost_priority": 1.98e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.96e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 1.1e-05, + "output_cost_per_token": 1.1e-05, + "output_cost_per_token_flex": 5.5e-06, + "output_cost_per_token_priority": 1.98e-05, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.2e-07, + "input_cost_per_token": 2.6e-06, + "input_cost_per_token_above_200k_tokens": 5.2e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.56e-05, + "output_cost_per_token_above_200k_tokens": 2.34e-05, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "input_cost_per_token_priority": 4.68e-06, + "input_cost_per_token_above_200k_tokens_priority": 9.36e-06, + "output_cost_per_token_priority": 2.808e-05, + "output_cost_per_token_above_200k_tokens_priority": 4.212e-05, + "cache_read_input_token_cost_priority": 4.68e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 9.36e-07, + "cache_read_input_token_cost_flex": 2.6e-07, + "input_cost_per_token_flex": 1.3e-06, + "output_cost_per_token_flex": 7.8e-06, + }, + "gemini/gemini-3.8-flash": { + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token": 5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 6.25e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_audio_token": 6e-06, "output_cost_per_reasoning_token": 3.5e-06, "output_cost_per_token": 3e-06, "output_cost_per_token_flex": 1.5e-06, @@ -454,26 +532,92 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_function_calling": True, "web_search_billing_unit": "per_query", }, - "gemini-3.8-flash": { - "cache_read_input_token_cost": 5.2e-08, - "google_maps_grounding_cost_per_query": 0.025, - "input_cost_per_audio_token": 1.04e-06, - "input_cost_per_token": 5.2e-07, - "input_cost_per_token_flex": 2.6e-07, - "input_cost_per_token_priority": 6.5e-07, - "input_cost_per_video_token": 6.2e-07, - "litellm_provider": "vertex_ai-language-models", + "vertex_ai/gemini-3.1-pro-preview": { + "cache_creation_input_token_cost_above_200k_tokens": 2.75e-07, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4.4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.92e-07, + "cache_read_input_token_cost_priority": 3.96e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_200k_tokens": 4.4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.92e-06, + "input_cost_per_token_flex": 1.1e-06, + "input_cost_per_token_priority": 3.96e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_image": 0.000132, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_200k_tokens": 1.98e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.564e-05, + "output_cost_per_token_flex": 6.6e-06, + "output_cost_per_token_priority": 2.376e-05, + "google_maps_grounding_cost_per_query": 0.0154, + "search_context_cost_per_query": { + "search_context_size_high": 0.0154, + "search_context_size_low": 0.0154, + "search_context_size_medium": 0.0154, + }, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + }, + "vertex_ai/gemini-3.8-flash": { + "cache_read_input_token_cost": 5.2e-08, + "cache_read_input_token_cost_flex": 2.6e-08, + "cache_read_input_token_cost_priority": 9.36e-08, + "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token": 5.2e-07, + "input_cost_per_token_flex": 2.6e-07, + "input_cost_per_token_priority": 6.5e-07, + "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_audio_token": 6.24e-06, "output_cost_per_token": 3.12e-06, + "output_cost_per_reasoning_token": 2.6e-06, "output_cost_per_token_flex": 1.56e-06, "output_cost_per_token_priority": 3.9e-06, - "search_context_cost_per_query": {"search_context_size_medium": 0.035}, + "search_context_cost_per_query": { + "search_context_size_high": 0.035, + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + }, "supports_function_calling": True, - "web_search_billing_unit": "per_prompt", + "supports_prompt_caching": True, + "supports_reasoning": True, + "web_search_billing_unit": "per_query", + }, + "vertex_ai/gemini-4-argon": { + "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_flex": 6e-08, + "cache_read_input_token_cost_priority": 2.16e-07, + "google_maps_grounding_cost_per_query": 0.0168, + "input_cost_per_audio_token": 2.4e-06, + "input_cost_per_token": 2.4e-06, + "input_cost_per_token_flex": 1.2e-06, + "input_cost_per_token_priority": 4.32e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_reasoning_token": 1.2e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.0168, + "search_context_size_low": 0.0168, + "search_context_size_medium": 0.0168, + }, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "web_search_billing_unit": "per_query", }, "together_ai/moonshotai/Kimi-K3": { "input_cost_per_token": 1.15e-06, @@ -620,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", @@ -967,6 +1117,126 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "supports_reasoning": True, "supports_prompt_caching": True, "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6.1-sol": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 2.42e-06, + "output_cost_per_token": 1.21e-05, + "cache_read_input_token_cost": 1.21e-07, + "cache_creation_input_token_cost": 3.025e-06, + "input_cost_per_token_above_272k_tokens": 4.84e-06, + "output_cost_per_token_above_272k_tokens": 1.815e-05, + "cache_read_input_token_cost_above_272k_tokens": 2.42e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.05e-06, + "input_cost_per_token_ultrafast": 1.452e-05, + "output_cost_per_token_ultrafast": 7.26e-05, + "cache_read_input_token_cost_ultrafast": 7.26e-07, + "cache_creation_input_token_cost_ultrafast": 1.815e-05, + "input_cost_per_token_above_272k_tokens_ultrafast": 2.904e-05, + "output_cost_per_token_above_272k_tokens_ultrafast": 1.089e-04, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.452e-06, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 3.63e-05, + "max_tokens": 131072, + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6-luna": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 1.32e-07, + "output_cost_per_token": 6.6e-07, + "cache_read_input_token_cost": 1.32e-08, + "cache_creation_input_token_cost": 1.65e-07, + "input_cost_per_token_above_272k_tokens": 2.64e-07, + "output_cost_per_token_above_272k_tokens": 9.9e-07, + "cache_read_input_token_cost_above_272k_tokens": 2.64e-08, + "cache_creation_input_token_cost_above_272k_tokens": 3.3e-07, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-6-astra": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 1.43e-05, + "output_cost_per_token": 7.15e-05, + "cache_read_input_token_cost": 1.43e-06, + "cache_creation_input_token_cost": 1.7875e-05, + "input_cost_per_token_above_272k_tokens": 2.86e-05, + "output_cost_per_token_above_272k_tokens": 1.0725e-04, + "cache_read_input_token_cost_above_272k_tokens": 2.86e-06, + "cache_creation_input_token_cost_above_272k_tokens": 3.575e-05, + "input_cost_per_token_ultrafast": 8.58e-05, + "output_cost_per_token_ultrafast": 4.29e-04, + "cache_read_input_token_cost_ultrafast": 8.58e-06, + "cache_creation_input_token_cost_ultrafast": 1.0725e-04, + "input_cost_per_token_above_272k_tokens_ultrafast": 1.716e-04, + "output_cost_per_token_above_272k_tokens_ultrafast": 6.435e-04, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.716e-05, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 2.145e-04, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_none_reasoning_effort": False, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + }, + "bedrock_mantle/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "input_cost_per_token": 3.08e-06, + "output_cost_per_token": 1.848e-05, + "cache_read_input_token_cost": 3.08e-07, + "cache_creation_input_token_cost": 3.85e-06, + "input_cost_per_token_above_272k_tokens": 6.16e-06, + "output_cost_per_token_above_272k_tokens": 2.772e-05, + "cache_read_input_token_cost_above_272k_tokens": 6.16e-07, + "cache_creation_input_token_cost_above_272k_tokens": 7.7e-06, + "max_tokens": 128000, + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_reasoning": True, + "supports_prompt_caching": True, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "use_openai_responses_path": True, + "search_context_cost_per_query": { + "search_context_size_high": 0.0168, + "search_context_size_low": 0.0168, + "search_context_size_medium": 0.0168, + }, + }, + "bedrock_mantle/anthropic.claude-opus-5-5": { + "litellm_provider": "bedrock_mantle", + "mode": "chat", + "input_cost_per_token": 6.6e-06, + "output_cost_per_token": 3.3e-05, + "cache_read_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost": 8.25e-06, + "cache_creation_input_token_cost_above_1hr": 1.32e-05, + "max_tokens": 128000, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "supports_function_calling": True, + "supports_reasoning": True, + "supports_prompt_caching": True, }, "us.anthropic.claude-opus-5-5": { "cache_creation_input_token_cost_above_1hr": 8.8e-06, @@ -1259,29 +1529,56 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token_above_272k_tokens_priority": 4.6e-05, "cache_read_input_token_cost_above_272k_tokens": 4.6e-07, }, + "azure/gpt-5.6-terra": { + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.75e-06, + "cache_creation_input_token_cost_priority": 5.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.1e-05, + "cache_creation_input_token_cost_flex": 1.375e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.2e-07, + "cache_read_input_token_cost_priority": 4.4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.8e-07, + "cache_read_input_token_cost_flex": 1.1e-07, + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2.2e-06, + "input_cost_per_token_priority": 4.4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.8e-06, + "input_cost_per_token_flex": 1.1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "output_cost_per_token_above_272k_tokens_flex": 9.9e-06, + "output_cost_per_token_priority": 2.64e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.96e-05, + "output_cost_per_token_flex": 6.6e-06, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + }, "azure/gpt-6-luna": { "cache_read_input_token_cost": 1.3e-08, "cache_creation_input_token_cost": 1.625e-07, "cache_creation_input_token_cost_above_272k_tokens": 3.25e-07, "input_cost_per_token": 1.3e-07, - "input_cost_per_token_flex": 6.5e-08, - "input_cost_per_token_priority": 2.6e-07, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.5e-07, - "output_cost_per_token_flex": 3.25e-07, - "output_cost_per_token_priority": 1.3e-06, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.6e-07, - "input_cost_per_token_above_272k_tokens_flex": 1.3e-07, - "input_cost_per_token_above_272k_tokens_priority": 5.2e-07, "output_cost_per_token_above_272k_tokens": 1.3e-06, - "output_cost_per_token_above_272k_tokens_flex": 6.5e-07, - "output_cost_per_token_above_272k_tokens_priority": 2.6e-06, "cache_read_input_token_cost_above_272k_tokens": 2.6e-08, }, "azure/gpt-6-astra": { @@ -1289,24 +1586,16 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_creation_input_token_cost": 1.5625e-05, "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, "input_cost_per_token": 1.25e-05, - "input_cost_per_token_flex": 6.25e-06, - "input_cost_per_token_priority": 2.5e-05, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.25e-05, - "output_cost_per_token_flex": 3.125e-05, - "output_cost_per_token_priority": 0.000125, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 2.5e-05, - "input_cost_per_token_above_272k_tokens_flex": 1.25e-05, - "input_cost_per_token_above_272k_tokens_priority": 5e-05, "output_cost_per_token_above_272k_tokens": 0.000125, - "output_cost_per_token_above_272k_tokens_flex": 6.25e-05, - "output_cost_per_token_above_272k_tokens_priority": 0.00025, "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, }, "azure/gpt-6.1-sol": { @@ -1314,24 +1603,16 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 2.4e-06, - "input_cost_per_token_flex": 1.2e-06, - "input_cost_per_token_priority": 4.8e-06, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.2e-05, - "output_cost_per_token_flex": 6e-06, - "output_cost_per_token_priority": 2.4e-05, "supports_function_calling": True, "supports_none_reasoning_effort": True, "input_cost_per_token_above_272k_tokens": 4.8e-06, - "input_cost_per_token_above_272k_tokens_flex": 2.4e-06, - "input_cost_per_token_above_272k_tokens_priority": 9.6e-06, "output_cost_per_token_above_272k_tokens": 2.4e-05, - "output_cost_per_token_above_272k_tokens_flex": 1.2e-05, - "output_cost_per_token_above_272k_tokens_priority": 4.8e-05, "cache_read_input_token_cost_above_272k_tokens": 4.8e-07, }, } diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 038c8d6b9e3..3efa97407f7 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -48,6 +48,7 @@ class CostMapEntry(BaseModel): supports_reasoning: bool | None = None supports_prompt_caching: bool | None = None supported_endpoints: tuple[str, ...] | None = None + use_openai_responses_path: bool | None = None supports_bedrock_runtime_chat_completions_tools_with_reasoning: bool | None = None supports_bedrock_runtime_chat_completions_response_format: bool | None = None cache_creation_input_token_cost_above_100k_tokens: float | None = None @@ -235,6 +236,8 @@ class RealtimeResponse(BaseModel): session_type: str | None = None created_event: Literal["session.created", "transcription_session.created"] = "session.created" created_repeats: int = 1 + created_after_first_frame: bool = False + dialect: Literal["openai", "xai"] = "openai" StoredResponse: TypeAlias = Annotated[ @@ -570,6 +573,7 @@ _PROVIDER_PREFIXES: Final[Mapping[str, str]] = MappingProxyType( "deepgram": "deepgram", "text-completion-openai": "text-completion-openai", "cohere": "cohere", + "vertex_ai": "vertex_ai", "vertex_ai-language-models": "vertex_ai", "vertex_ai-image-models": "vertex_ai", "vertex_ai-anthropic_models": "vertex_ai", @@ -619,6 +623,9 @@ _LITELLM_PARAMS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( "deepgram": MappingProxyType({}), "text-completion-openai": MappingProxyType({}), "cohere": MappingProxyType({}), + "vertex_ai": MappingProxyType( + {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} + ), "vertex_ai-language-models": MappingProxyType( {"vertex_project": "cc-scripted-project", "vertex_location": "us-central1"} ), diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index c8de39c0d38..60bd7ae72a3 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -106,6 +106,82 @@ def test_runner_environment_drops_inherited_provider_api_keys() -> None: assert inherited["OPENAI_API_KEY"] == "ambient-openai" +@pytest.mark.parametrize("fails", (False, True), ids=("passing-test", "failing-test")) +def test_sdk_teardown_delivers_logging_before_restoring_callbacks(tmp_path: Path, fails: bool) -> None: + root: Final = Path(__file__).resolve().parents[3] + events: Final = tmp_path / "callbacks.json" + probe: Final = tmp_path / "test_callback_teardown.py" + probe.write_text( + "import asyncio, json, litellm, pytest\n" + "from pathlib import Path\n" + "from litellm.integrations.custom_logger import CustomLogger\n" + f"EVENTS = Path({str(events)!r})\n" + "ORIGINAL = tuple(litellm.callbacks)\n" + "@pytest.fixture\n" + "def release():\n" + " gate = asyncio.Event()\n" + " yield gate\n" + " gate.set()\n" + "async def test_request(release):\n" + " started = asyncio.Event()\n" + " request_loop = asyncio.get_running_loop()\n" + " class Recorder(CustomLogger):\n" + " async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):\n" + " started.set()\n" + " await release.wait()\n" + " EVENTS.write_text(json.dumps({\n" + " 'registered': self in litellm.callbacks,\n" + " 'same_loop': asyncio.get_running_loop() is request_loop,\n" + " 'call_id': kwargs['litellm_call_id'],\n" + " }))\n" + " litellm.callbacks.append(Recorder())\n" + " response = await litellm.acompletion(model='gpt-4o-mini',\n" + " messages=[{'role': 'user', 'content': 'teardown'}],\n" + " mock_response='done', litellm_call_id='teardown-request')\n" + " await asyncio.wait_for(started.wait(), timeout=5)\n" + " assert response.choices[0].message.content == 'done'\n" + " assert not EVENTS.exists()\n" + f" assert not {fails}, 'intentional test failure'\n" + "def test_next():\n" + " assert EVENTS.exists(), 'logging did not finish before the next test'\n" + " assert json.loads(EVENTS.read_text()) == {\n" + " 'registered': True, 'same_loop': True, 'call_id': 'teardown-request'}\n" + " assert tuple(litellm.callbacks) == ORIGINAL\n" + ) + result: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "pytest", + str(probe), + "--noconftest", + "-p", + "tests.integration.sdk.conftest", + "-o", + "addopts=", + "-o", + "asyncio_mode=auto", + "-o", + "asyncio_default_fixture_loop_scope=session", + f"--junitxml={tmp_path / 'junit.xml'}", + "-q", + ), + cwd=root, + env=_nested_pytest_environment(os.environ), + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == (1 if fails else 0), result.stdout + result.stderr + cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase") + assert tuple(case.attrib["name"] for case in cases) == ("test_request", "test_next") + assert all(case.find("error") is None for case in cases), result.stdout + result.stderr + assert (cases[0].find("failure") is not None) == fails, result.stdout + result.stderr + assert cases[1].find("failure") is None, result.stdout + result.stderr + + def test_owned_readiness_rejects_a_healthy_listener_from_another_process( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: 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/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c3bccec1f8e..e288cef3699 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -47,10 +47,20 @@ pytestmark: Final = pytest.mark.xdist_group("mcp_llm_private_gateway") @pytest.fixture(scope="session") def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway") + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = directory / "stored-responses.yaml" + config.write_text( + yaml.safe_dump( + { + **stock, + "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, + } + ) + ) with ( worker_services(directory) as services, gateway_from_environment() as shared, - owned_proxy(shared, directory, services) as private, + owned_proxy(shared, directory, services, config=config) as private, ): yield private, services @@ -106,6 +116,17 @@ class Turn: answer: str +@dataclass(frozen=True, slots=True) +class Billing: + first: tuple[int, int] = (10, 5) + follow_up: tuple[int, int] = (10, 5) + follow_up_fails: bool = False + + +EVEN_ROUNDS: Final = Billing() +FOLLOW_UP_OUTAGE: Final = "scripted follow-up outage" + + def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]: return lambda _: Turn(tool, json.dumps(ADD), ANSWER) @@ -162,7 +183,58 @@ def _responses_stream(response: Mapping[str, JsonValue], item: Mapping[str, Json ) -def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[Request], Reply]: +def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]: + content: Final = message.get("content") if isinstance(message, dict) else None + blocks: Final = content if isinstance(content, list) else [] + return frozenset(str(block[field]) for block in blocks if isinstance(block, dict) and block.get("type") == kind) + + +def _unanswered_tool_uses(body: Mapping[str, JsonValue]) -> frozenset[str]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + followers: Final = [*turns[1:], None] + return frozenset().union( + *( + _block_ids(turn, "tool_use", "id") - _block_ids(follower, "tool_result", "tool_use_id") + for turn, follower in zip(turns, followers) + ) + ) + + +def _orphan_tool_results(body: Mapping[str, JsonValue]) -> frozenset[str]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + leaders: Final = [None, *turns[:-1]] + return frozenset().union( + *( + _block_ids(turn, "tool_result", "tool_use_id") - _block_ids(leader, "tool_use", "id") + for turn, leader in zip(turns, leaders) + ) + ) + + +def _anthropic_invalid_request(message: str) -> Reply: + error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": message}} + return Reply(status=400, body=json.dumps(error).encode()) + + +def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None: + unanswered: Final = _unanswered_tool_uses(body) + orphans: Final = _orphan_tool_results(body) + if unanswered: + return _anthropic_invalid_request( + f"`tool_use` ids were found without `tool_result` blocks immediately after: {', '.join(sorted(unanswered))}" + ) + if orphans: + return _anthropic_invalid_request( + f"unexpected `tool_use_id` found in `tool_result` blocks: {', '.join(sorted(orphans))}" + ) + return None + + +def _model_double( + plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS +) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: if request.method == "GET" and request.target.endswith("/models"): return _json({"object": "list", "data": []}) @@ -172,8 +244,16 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ return Reply(status=500, body=json.dumps(outage).encode()) turn: Final = plan(body) done: Final = _has_tool_result(body) + if done and billing.follow_up_fails: + failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}} + return Reply(status=500, body=json.dumps(failure).encode()) + prompt_tokens, completion_tokens = billing.follow_up if done else billing.first identity: Final = uuid.uuid4().hex[:12] - usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + usage: Final = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } if request.target.endswith("/chat/completions"): message: Final = ( {"role": "assistant", "content": turn.answer} @@ -222,6 +302,8 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ } ) if request.target.endswith("/messages"): + if (rejected := _anthropic_tool_result_rule(body)) is not None: + return rejected content: Final = ( [{"type": "text", "text": turn.answer}] if done @@ -236,7 +318,7 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ "content": content, "stop_reason": "end_turn" if done else "tool_use", "stop_sequence": None, - "usage": {"input_tokens": 10, "output_tokens": 5}, + "usage": {"input_tokens": prompt_tokens, "output_tokens": completion_tokens}, } ) assert request.target.endswith("/responses"), request.target @@ -265,7 +347,11 @@ def _model_double(plan: Callable[[Mapping[str, JsonValue]], Turn]) -> Callable[[ "status": "completed", "model": "gpt-4o-mini", "output": [item], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "usage": { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, } return _responses_stream(response, item) if body.get("stream") is True else _json(response) @@ -354,16 +440,28 @@ def _upstream_model(surface: Surface) -> str: }[surface] +PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "num_retries": 0} + + @contextmanager -def _rig(gateway: Gateway, surface: Surface) -> Iterator[Rig]: +def _rig( + gateway: Gateway, + surface: Surface, + *, + upstream: Surface | None = None, + billing: Billing = EVEN_ROUNDS, + prices: Mapping[str, JsonValue] | None = None, +) -> Iterator[Rig]: alias: Final = "llm" + uuid.uuid4().hex[:8] with ( mcp_peer() as peer, - wire_server(_model_double(_fixed_turn(f"{alias}-add"))) as wire, + wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing)) as wire, gateway.scenario() as scenario, ): server_id: Final = register_mcp(scenario, peer, alias) - model: Final = scenario.model(model=_upstream_model(surface), api_base=wire.url + "/v1") + model: Final = scenario.model( + model=_upstream_model(upstream or surface), api_base=wire.url + "/v1", **(prices or {}) + ) peer.drain() yield Rig(gateway, scenario, peer, wire, alias, server_id, model, surface) @@ -590,6 +688,119 @@ def test_an_open_server_in_another_gateway_cannot_widen_this_gateways_tools(gate assert deleted.json()["info"]["status"] == "deleted", deleted.text +FIRST_ROUND: Final = (11, 3) +FOLLOW_UP_ROUND: Final = (17, 7) + + +def _price(*rounds: tuple[int, int]) -> float: + input_price, output_price = PRICES["input_cost_per_token"], PRICES["output_cost_per_token"] + assert isinstance(input_price, float) and isinstance(output_price, float) + return sum(prompt * input_price + completion * output_price for prompt, completion in rounds) + + +def _responses_rows(call_id: str, status: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + "SELECT request_id, spend, prompt_tokens, completion_tokens, session_id " + 'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "aresponses", status), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + + +def _session_call_types(session_id: str) -> frozenset[str]: + rows: Final = read_rows('SELECT call_type FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (session_id,)) + return frozenset(str(row["call_type"]) for row in rows) + + +def _message_text(message: Mapping[str, JsonValue]) -> str: + content: Final = message.get("content") + if isinstance(content, str): + return content + blocks: Final = content if isinstance(content, list) else [] + return "".join(str(block["text"]) for block in blocks if isinstance(block, dict) and block.get("type") == "text") + + +def _assistant_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + messages: Final = body.get("messages") + turns: Final = messages if isinstance(messages, list) else [] + return tuple(_message_text(turn) for turn in turns if isinstance(turn, dict) and turn.get("role") == "assistant") + + +def test_nonstream_responses_auto_execute_logs_the_answer_and_bills_both_rounds(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND) + with _rig(gateway, "responses", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert rig.final_text(body) == ANSWER, response.text + rows: Final = _responses_rows(response.headers["x-litellm-call-id"], "success") + assert len(rows) == 1, rows + row: Final = rows[0] + assert row["request_id"] == body["id"], row + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + FIRST_ROUND[0] + FOLLOW_UP_ROUND[0], + FIRST_ROUND[1] + FOLLOW_UP_ROUND[1], + ), row + assert object_value(body["usage"])["input_tokens"] == row["prompt_tokens"], body + assert float(str(row["spend"])) == pytest.approx(_price(FIRST_ROUND, FOLLOW_UP_ROUND)), row + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(float(str(row["spend"]))) + session_id: Final = string_value(row["session_id"]) + eventually( + lambda: _session_call_types(session_id), + lambda found: {"list_mcp_tools", "call_mcp_tool"} <= found, + seconds=70, + ) + + +def test_nonstream_responses_failed_follow_up_bills_round_one_on_the_failure_row(gateway: Gateway) -> None: + billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True) + with _rig(gateway, "responses", billing=billing, prices=PRICES) as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 500 and FOLLOW_UP_OUTAGE in response.text, response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = _responses_rows(call_id, "failure") + assert len(rows) == 1, rows + assert float(str(rows[0]["spend"])) == pytest.approx(_price(FIRST_ROUND)), rows + assert ( + read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s AND status=%s', + (call_id, "aresponses", "success"), + ) + == [] + ) + + +def test_bridged_responses_auto_execute_answers_every_tool_use_with_its_result(gateway: Gateway) -> None: + with _rig(gateway, "responses", upstream="messages") as rig: + response: Final = rig.send(_granted_key(rig), [AUTO]) + assert response.status_code == 200, response.text + assert rig.final_text(response.json()) == ANSWER, response.text + assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] + + +def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None: + with _rig(gateway, "responses", upstream="messages") as rig: + key: Final = _granted_key(rig) + first: Final = rig.send(key, [AUTO]) + assert first.status_code == 200, first.text + first_id: Final = string_value(JSON_OBJECT.validate_json(first.content)["id"]) + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (first_id,)), + lambda found: len(found) == 1, + seconds=70, + ) + rig.wire.drain() + second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True}) + assert second.status_code == 200, second.text + opening: Final = next( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert ANSWER in _assistant_texts(opening), opening + + Bridge = Literal["chat", "responses", "messages"] BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses") Client = Literal["sync", "async"] 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_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 605130acd38..2ca9e0c6410 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -577,7 +577,6 @@ def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gat def test_two_identical_cache_hit_rejects_write_matching_spend_rows(gateway: Gateway, tmp_path: Path) -> None: """E1b: both cache-hit rejects land a failure spend row.""" - pytest.skip("BUG: roughly one in four back-to-back cache-hit rejects never lands its LiteLLM_SpendLogs row") marker: Final = uuid.uuid4().hex with _rig(gateway, tmp_path, marker) as rig: text: Final = "cache hit control e1b " + marker diff --git a/tests/integration/observability/test_callback_delivery.py b/tests/integration/observability/test_callback_delivery.py index 543e60da4a6..e786e655a13 100644 --- a/tests/integration/observability/test_callback_delivery.py +++ b/tests/integration/observability/test_callback_delivery.py @@ -11,7 +11,7 @@ import pytest import yaml from integration._support.client import Gateway, JsonValue, eventually, object_value, string_value from integration._support.database import read_rows -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 @@ -268,6 +268,89 @@ def test_streamed_responses_success_callback_carries_provider_apim_request_id(ga }, json.dumps(events[0]["hidden_params"]) +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_non_stream_openai_like_success_callback_carries_provider_headers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = uuid.uuid4().hex + request_id: Final = "req_" + marker + sink_secret: Final = "synthetic-sink-secret-" + marker + + def upstream(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + marker, + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode(), + headers={"x-amzn-requestid": request_id, "x-ratelimit-remaining-requests": "997"}, + ) + + def sink(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {sink_secret}" + return Reply() + + with wire_server(upstream) as provider, wire_server(sink) as endpoint: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy( + gateway, + tmp_path, + { + "GENERIC_LOGGER_ENDPOINT": endpoint.url, + "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}", + }, + config=path, + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model="bedrock_mantle/google.gemma-4-31b", + api_base=provider.url, + api_key="synthetic-mantle-bearer", + aws_region_name="us-west-2", + ) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "header callback " + marker}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == marker, response.text + assert len(provider.drain()) == 1 + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered() -> tuple[dict, ...]: + batches.extend(endpoint.drain()) + return tuple( + event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model + ) + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10) + assert (events[0]["status"], events[0]["call_type"]) == ("success", "acompletion") + assert events[0]["litellm_call_id"] == response.headers["x-litellm-call-id"] + additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {} + assert { + name: additional_headers.get(name) + for name in ("llm_provider-x-amzn-requestid", "x_ratelimit_remaining_requests") + } == {"llm_provider-x-amzn-requestid": request_id, "x_ratelimit_remaining_requests": 997}, json.dumps( + events[0]["hidden_params"] + ) + + _RAISING_HOOK: Final = """ from litellm.integrations.custom_logger import CustomLogger 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_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py index ccef839a163..b36da940a42 100644 --- a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -8,6 +8,11 @@ dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult documents the proxy forwarded to the peer. +The fallback cells cover a websocket that outlives one provider attempt. A refused stream keeps the +client's opening message for the fallback, which gets it once whether it dials Bedrock again or +an OpenAI realtime deployment, and once Bedrock accepted a stream no later attempt may open a second +provider session on that websocket. + The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. """ @@ -21,12 +26,19 @@ import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager from pathlib import Path +from types import MappingProxyType from typing import Final +import httpx import pytest import websockets +import yaml from integration._support.bedrock_realtime_peer import ( + BedrockEvent, BedrockRealtimePeer, + OpeningEnd, + Refusal, + StreamFailure, Turn, assistant_text, bedrock_realtime_peer, @@ -43,12 +55,15 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import write_self_signed_cert +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import RealtimeResponse from pydantic import JsonValue from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -60,6 +75,30 @@ WEATHER_FUNCTION: Final = { NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' +QUESTION: Final = "What is the weather in New York right now?" +ANSWER: Final = "It is 72 and sunny" +INSTRUCTIONS: Final = "Answer every weather question in one short sentence." +OPENAI_VOICE_MODEL: Final = "gpt-realtime-2.1" +BEDROCK_REPLAY_GROUPS: Final = MappingProxyType( + { + "system_prompt": ("sonic-refused-session-update-first", "sonic-replay-session-update-first"), + "user_input": ("sonic-refused-text-first", "sonic-replay-text-first"), + } +) +OPENAI_REPLAY_GROUPS: Final = ("sonic-refused-openai-fallback", "openai-replay-target") +OPENAI_COMMITTED_GROUPS: Final = ("sonic-committed-openai-fallback", "openai-committed-target") +FALLBACK_GROUPS: Final = (*BEDROCK_REPLAY_GROUPS.values(), OPENAI_REPLAY_GROUPS, OPENAI_COMMITTED_GROUPS) +SCRIPTED_DONE: Final[dict[str, JsonValue]] = { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": { + "id": "resp_$UNIQUE_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, + }, +} @pytest.fixture(scope="module") @@ -67,15 +106,35 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) +def _fallback_config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + config: Final = { + **stock, + "router_settings": { + **object_value(stock["router_settings"]), + "fallbacks": [{source: [target]} for source, target in FALLBACK_GROUPS], + }, + } + path: Final = directory / "bedrock-sonic-fallbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("bedrock-sonic") with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_fallback_config(directory)) as proxy: yield proxy def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: - name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + return _bedrock_deployment(gateway, scenario, peer, f"bedrock-sonic-{uuid.uuid4().hex[:12]}") + + +def _bedrock_deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer, name: str) -> str: created: Final = gateway.post( "/model/new", { @@ -94,6 +153,44 @@ def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) return name +def _openai_deployment(gateway: Gateway, scenario: Scenario, name: str) -> str: + scenario_id: Final = f"bedrock-sonic-fallback-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario( + scenario_id, RealtimeResponse(content_type="application/x-realtime", events=(SCRIPTED_DONE,)) + ) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"openai/{OPENAI_VOICE_MODEL}", + "api_key": scenario_id, + "api_base": f"{gateway.upstream_url}/v1", + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return scenario_id + + +def _upstream_requests(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content) + requests: Final = observations["requests"] + assert isinstance(requests, list), observations + return tuple( + request + for request in map(object_value, requests) + if request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id + ) + + +def _upstream_frames(requests: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(request["body"]) for request in requests if request["method"] == "WEBSOCKET_FRAME") + + def _session_update(tool: Mapping[str, JsonValue]) -> str: return json.dumps( { @@ -113,12 +210,16 @@ def _user_text() -> str: "item": { "type": "message", "role": "user", - "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + "content": [{"type": "input_text", "text": QUESTION}], }, } ) +def _instructions_update() -> str: + return json.dumps({"type": "session.update", "session": {"instructions": INSTRUCTIONS}}) + + _RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) @@ -207,6 +308,7 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple assert _tail_types(events) == ( "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.done", @@ -376,3 +478,160 @@ def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], con events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) (call,) = _function_calls(events) assert string_value(call["arguments"]) == (content if content is not None else "{}"), events + + +def _of_type(events: tuple[Event, ...], event_type: str) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == event_type) + + +def _bedrock_fields(events: tuple[BedrockEvent, ...], name: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(object_value(event["event"]) for event in events) + return tuple(object_value(body[name]) for body in bodies if name in body) + + +def _text_inputs(events: tuple[BedrockEvent, ...], role: str) -> tuple[str, ...]: + roles: Final = MappingProxyType( + {string_value(start["contentName"]): start.get("role") for start in _bedrock_fields(events, "contentStart")} + ) + return tuple( + string_value(text["content"]) + for text in _bedrock_fields(events, "textInput") + if roles.get(string_value(text["contentName"])) == role + ) + + +def _answer_turn() -> Turn: + return Turn("user_input", assistant_text(ANSWER) + (usage(30, 12),)) + + +def _answer_text(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in _of_type(events, "response.text.delta")) + + +async def _opened_with(socket: ClientConnection, opening: tuple[str, ...], stop: frozenset[str]) -> tuple[Event, ...]: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + for message in opening: + await socket.send(message) + return (created,) + await _events_until(socket, lambda event, _events: event.get("type") in stop) + + +async def _session( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...], stop: frozenset[str] +) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + return await _opened_with(socket, opening, stop) + + +async def _session_until_closed( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...] +) -> tuple[tuple[Event, ...], int | None]: + async with _realtime(gateway, model, key) as socket: + events: Final = await _opened_with(socket, opening, frozenset(("error", "session.created"))) + if events[-1].get("type") != "error": + return events, None + try: + await asyncio.wait_for(socket.recv(), 30) + except ConnectionClosed: + return events, socket.close_code + return events, None + + +@pytest.mark.parametrize("after", ["system_prompt", "user_input"], ids=["session-update-first", "text-first"]) +def test_refused_opening_is_replayed_on_the_bedrock_fallback( + candidate: Gateway, cert: tuple[Path, Path], after: OpeningEnd +) -> None: + source, target = BEDROCK_REPLAY_GROUPS[after] + update_first: Final = after == "system_prompt" + opening: Final = ((_instructions_update(),) if update_first else ()) + (_user_text(), _RESPONSE_CREATE) + opening_role, opening_text = ("SYSTEM", INSTRUCTIONS) if update_first else ("USER", QUESTION) + with ( + bedrock_realtime_peer((), cert, refusal=Refusal(after)) as refusing, + bedrock_realtime_peer((_answer_turn(),), cert) as healthy, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + _bedrock_deployment(candidate, scenario, healthy, target) + events: Final = asyncio.run( + _session( + candidate, source, scenario.key(), opening, frozenset(("response.done", "error", "session.created")) + ) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_of_type(events, "session.updated")) == int(update_first), events + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + (replayed,) = healthy.streams() + assert _text_inputs(replayed, "USER") == (QUESTION,), replayed + assert _text_inputs(replayed, "SYSTEM").count(INSTRUCTIONS) == int(update_first), replayed + refused: Final = refusing.streams() + assert refused and all(opening_text in _text_inputs(stream, opening_role) for stream in refused), refused + + +def test_refused_text_opening_is_replayed_on_the_openai_fallback(candidate: Gateway, cert: tuple[Path, Path]) -> None: + source, target = OPENAI_REPLAY_GROUPS + with ( + bedrock_realtime_peer((), cert, refusal=Refusal("user_input")) as refusing, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events: Final = asyncio.run( + _session( + candidate, + source, + scenario.key(), + (_user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "error")), + ) + ) + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + requests: Final = _upstream_requests(candidate, scenario_id) + assert len(tuple(request for request in requests if request["method"] == "WEBSOCKET")) == 1, requests + assert _upstream_frames(requests) == (json.loads(_user_text()), json.loads(_RESPONSE_CREATE)), requests + + +def test_committed_session_failure_closes_the_socket_without_another_provider_stream( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + source, target = OPENAI_COMMITTED_GROUPS + with ( + bedrock_realtime_peer((_answer_turn(),), cert, failure=StreamFailure()) as committed, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, committed, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events, close_code = asyncio.run( + _session_until_closed(candidate, source, scenario.key(), (_user_text(), _RESPONSE_CREATE)) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + assert string_value(events[-1]["type"]) == "error", events + assert close_code == 1011, events + assert _upstream_requests(candidate, scenario_id) == () + assert len(committed.streams()) == 1 + + +def test_voice_object_is_refused_and_never_reaches_bedrock(candidate: Gateway, cert: tuple[Path, Path]) -> None: + voice_update: Final = json.dumps({"type": "session.update", "session": {"voice": {"id": "voice_123"}}}) + with bedrock_realtime_peer((_answer_turn(),), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + events: Final = asyncio.run( + _session( + candidate, + model, + scenario.key(), + (voice_update, _user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "session.created")), + ) + ) + errors: Final = tuple(object_value(event["error"]) for event in _of_type(events, "error")) + assert tuple((error.get("code"), error.get("param")) for error in errors) == ( + ("invalid_type", "session.voice"), + ), events + assert _of_type(events, "session.updated") == (), events + assert len(_of_type(events, "session.created")) == 1, events + assert _answer_text(events) == ANSWER, events + assert "voice_123" not in json.dumps(peer.received()) 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_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 87cd3105cc7..44dc98c9d01 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -83,6 +83,7 @@ _PROVIDERS: Final = ( "perplexity/pplx-decider-v1-27b", ), _Provider("typesafe", "typesafe/jev-1.13.0", "/v1/systemone", "jev-1.13.0", _API_KEY, False, "typesafe/jev-1.13.0"), + _Provider("nace", "nace/drex-v2.0", "/v1/systemone", "drex-v2.0", _API_KEY, False, "nace/drex-v2.0"), _Provider( "openrouter", "openrouter/typesafe/jev-1.13", @@ -135,7 +136,7 @@ _PROVIDERS: Final = ( ), ) _PERPLEXITY: Final = _PROVIDERS[0] -_OPENROUTER: Final = _PROVIDERS[2] +_OPENROUTER: Final = next(provider for provider in _PROVIDERS if provider.name == "openrouter") _OPENROUTER_CHAT_MODEL: Final = "openrouter/openai/gpt-5-mini" _UNSUPPORTED_PROVIDER_MODEL: Final = "anthropic/claude-opus-5-5" _CONNECTION_ERROR: Final = "litellm.APIConnectionError" @@ -581,14 +582,16 @@ async def test_sdk_openrouter_connection_failures_raise_a_connection_error(gatew api_base=dead_api_base, api_key=_API_KEY, ) - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.APIConnectionError) as decisions_error: litellm.decisions( model=_OPENROUTER.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=dead_api_base, api_key=_API_KEY ) - with pytest.raises(litellm.APIConnectionError): + assert decisions_error.value.llm_provider == "openrouter" + with pytest.raises(litellm.APIConnectionError) as systemone_error: await litellm.asystemone( model=_OPENROUTER.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=dead_api_base, api_key=_API_KEY ) + assert systemone_error.value.llm_provider == "openrouter" def test_a_deployment_whose_provider_has_no_decisions_support_is_refused_naming_every_supported_provider( 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_openai_like_response_headers_wire.py b/tests/integration/providers/test_openai_like_response_headers_wire.py new file mode 100644 index 00000000000..1ffa77cf326 --- /dev/null +++ b/tests/integration/providers/test_openai_like_response_headers_wire.py @@ -0,0 +1,435 @@ +import itertools +import json +import math +import uuid +from collections.abc import Callable, Mapping, Sequence +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.prometheus_series import Sample, scrape, series_cap_config +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Deploy: TypeAlias = Callable[[Gateway, Scenario, Wire], str] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PROJECT: Final = "scripted-project" +_REPLY_TEXT: Final = "header probe reply" +_USAGE: Final[JsonValue] = {"prompt_tokens": 13, "completion_tokens": 4, "total_tokens": 17} +_UPSTREAM_COST: Final = "0.0421" +_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s' +_KEY_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s' +_PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} +) +_RATE_LIMIT_HEADERS: Final[Mapping[str, str]] = MappingProxyType( + {"x-ratelimit-remaining-requests": "0", "x-ratelimit-remaining-tokens": "856, abc"} +) +_METERED_CALLS: Final[tuple[tuple[str, bool], ...]] = ( + ("groq", False), + ("groq", True), + ("openai", False), + ("openai", True), +) +_LATIN1_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( + ("chat", "bedrock_mantle", False), + ("chat", "bedrock_mantle", True), + ("responses", "groq", False), + ("responses", "groq", True), + ("messages", "groq", False), + ("messages", "groq", True), +) + + +def _keyed(model: str, **extra: JsonValue) -> Deploy: + def create(_gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model(model=model, api_base=wire.url, api_key="synthetic-provider-key", **extra) + + return create + + +def _vertex(model: str, location: str) -> Deploy: + def create(gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"vertex_ai/{model}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=location, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + return create + + +_MEMBERS: Final[Mapping[str, Deploy]] = MappingProxyType( + { + "bedrock_mantle": _keyed("bedrock_mantle/google.gemma-4-31b", aws_region_name="us-west-2"), + "groq": _keyed("groq/openai/gpt-oss-20b"), + "openai": _keyed("openai/gpt-5.4-mini"), + "watsonx": _keyed("watsonx/ibm/granite-3-8b-instruct", zen_api_key="synthetic-zen-key", project_id=_PROJECT), + "vertex_model_garden": _vertex("xai/grok-4.20-non-reasoning", "global"), + "maritalk": _keyed("maritalk/sabia-3"), + } +) + + +def _reply_json() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": "scripted", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": "stop"} + ], + "usage": _USAGE, + "service_tier": "on_demand", + } + ).encode() + + +def _chunk(response_id: str, choices: list[JsonValue], **extra: JsonValue) -> bytes: + body: Final = { + "id": response_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "scripted", + "choices": choices, + **extra, + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def _frames() -> tuple[bytes, ...]: + response_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return ( + _chunk( + response_id, [{"index": 0, "delta": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": None}] + ), + _chunk(response_id, [{"index": 0, "delta": {}, "finish_reason": "stop"}]), + _chunk(response_id, [], usage=_USAGE), + b"data: [DONE]\n\n", + ) + + +def _peer(headers: Mapping[str, str]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_frames(), headers=headers) + return Reply(body=_reply_json(), headers=headers) + + return respond + + +def _request_id_headers(request_id: str) -> Mapping[str, str]: + return MappingProxyType({"x-request-id": request_id, "x-amzn-requestid": request_id}) + + +def _body(endpoint: Endpoint, model: str, stream: bool) -> dict[str, JsonValue]: + prompt: Final = f"header probe {uuid.uuid4().hex}" + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + "cache": {"no-cache": True}, + **({"stream_options": {"include_usage": True}} if stream else {}), + } + case "responses": + return {"model": model, "input": prompt, "stream": stream} + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + } + + +def _send(gateway: Gateway, endpoint: Endpoint, model: str, key: str, *, stream: bool = False) -> httpx.Response: + return gateway.request("POST", _PATHS[endpoint], _body(endpoint, model, stream), key=key) + + +class _Reply(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Text(_Reply): + text: str = "" + + +class _ChatMessage(_Reply): + content: str = "" + + +class _ChatChoice(_Reply): + message: _ChatMessage + + +class _ChatReply(_Reply): + choices: tuple[_ChatChoice, ...] + + +class _ResponsesItem(_Reply): + content: tuple[_Text, ...] = () + + +class _ResponsesReply(_Reply): + output: tuple[_ResponsesItem, ...] + + +class _MessagesReply(_Reply): + content: tuple[_Text, ...] + + +def _text(endpoint: Endpoint, response: httpx.Response) -> str: + match endpoint: + case "chat": + return _ChatReply.model_validate_json(response.content).choices[0].message.content + case "responses": + items: Final = _ResponsesReply.model_validate_json(response.content).output + return "".join(part.text for part in itertools.chain.from_iterable(item.content for item in items)) + case "messages": + return "".join(block.text for block in _MessagesReply.model_validate_json(response.content).content) + + +def _header_cost(response: httpx.Response) -> float: + values: Final = response.headers.get_list("x-litellm-response-cost") + return float(values[0]) if values else math.nan + + +def _llm_provider_request_ids(response: httpx.Response) -> dict[str, str | None]: + return {name: response.headers.get(name) for name in ("llm_provider-x-request-id", "llm_provider-x-amzn-requestid")} + + +def _spend(key: str) -> float: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda rows: len(rows) >= 1, seconds=70) + rows: Final = held(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda found: len(found) == 1, holding=3) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float)), rows + return float(spend) + + +def _model_values(samples: Sequence[Sample], name: str, model: str) -> tuple[float, ...]: + return tuple( + sample.value for sample in samples if sample.name == name and sample.labels.get("model_group") == model + ) + + +def _deployment_metrics( + samples: Sequence[Sample], models: Mapping[str, str] +) -> dict[str, tuple[float, tuple[float, ...], tuple[float, ...]]]: + return { + call: ( + sum(_model_values(samples, "litellm_deployment_success_responses_total", model)), + _model_values(samples, "litellm_remaining_requests_metric", model), + _model_values(samples, "litellm_remaining_tokens_metric", model), + ) + for call, model in models.items() + } + + +def _key_spend(key: str) -> float | None: + rows: Final = read_rows(_KEY_SPEND_SELECT, (sha256(key.encode()).hexdigest(),)) + spend: Final = rows[0]["spend"] if rows else None + return float(spend) if isinstance(spend, (int, float)) else None + + +@pytest.mark.parametrize( + ("endpoint", "member"), + [ + ("chat", "bedrock_mantle"), + ("chat", "groq"), + ("chat", "watsonx"), + ("chat", "vertex_model_garden"), + ("chat", "maritalk"), + ("responses", "groq"), + ], +) +def test_non_stream_returns_the_provider_request_id(gateway: Gateway, endpoint: Endpoint, member: str) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + response: Final = _send(gateway, endpoint, model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert _text(endpoint, response) == _REPLY_TEXT, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_stream_chat_returns_the_provider_request_id(gateway: Gateway) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["bedrock_mantle"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model]), stream=True) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_upstream_cannot_overwrite_litellm_owned_headers(gateway: Gateway) -> None: + spoofed: Final = { + "x-litellm-call-id": "spoofed-call-id", + "x-litellm-model-id": "spoofed-model-id", + "x-litellm-attempted-fallbacks": "7", + } + with wire_server(_peer(spoofed)) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["groq"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-call-id"] != spoofed["x-litellm-call-id"], dict(response.headers) + assert response.headers["x-litellm-model-id"] != spoofed["x-litellm-model-id"], dict(response.headers) + assert response.headers["x-litellm-attempted-fallbacks"] == "0", dict(response.headers) + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize(("endpoint", "member"), [("chat", "bedrock_mantle"), ("responses", "groq")]) +def test_non_stream_spend_matches_stream_spend_under_an_upstream_cost_header( + gateway: Gateway, endpoint: Endpoint, member: str +) -> None: + with wire_server(_peer({"x-litellm-response-cost": _UPSTREAM_COST})) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + stream_key: Final = scenario.key(models=[model]) + non_stream_key: Final = scenario.key(models=[model]) + streamed: Final = _send(gateway, endpoint, model, stream_key, stream=True) + assert streamed.status_code == 200, streamed.text + assert _REPLY_TEXT in streamed.text, streamed.text + whole: Final = _send(gateway, endpoint, model, non_stream_key) + assert whole.status_code == 200, whole.text + assert _text(endpoint, whole) == _REPLY_TEXT, whole.text + stream_spend: Final = _spend(stream_key) + assert stream_spend > 0 + non_stream_spend: Final = _spend(non_stream_key) + assert math.isclose(non_stream_spend, stream_spend), (non_stream_spend, stream_spend) + assert len(wire.drain()) == 2 + + +@pytest.mark.parametrize("member", ["openai", "groq"]) +@pytest.mark.parametrize("upstream_cost", ["-5", "abc", "nan", "inf"], ids=["negative", "garbage", "nan", "infinite"]) +def test_an_invalid_upstream_cost_header_bills_the_price_map(gateway: Gateway, member: str, upstream_cost: str) -> None: + with ( + wire_server(_peer({"x-litellm-response-cost": upstream_cost})) as invalid_wire, + wire_server(_peer({})) as plain_wire, + gateway.scenario() as scenario, + ): + invalid_model: Final = _MEMBERS[member](gateway, scenario, invalid_wire) + plain_model: Final = _MEMBERS[member](gateway, scenario, plain_wire) + budget_key: Final = scenario.key(models=[invalid_model], max_budget=1) + reference_key: Final = scenario.key(models=[plain_model]) + billed: Final = _send(gateway, "chat", invalid_model, budget_key) + reference: Final = _send(gateway, "chat", plain_model, reference_key) + assert (billed.status_code, reference.status_code) == (200, 200), (billed.text, reference.text) + reference_spend: Final = _spend(reference_key) + assert reference_spend > 0 + key_spend: Final = eventually( + lambda: _key_spend(budget_key), + lambda spend: spend is not None and math.isclose(spend, reference_spend), + 70, + return_last_on_timeout=True, + ) + billed_costs: Final = { + "response_header": _header_cost(billed), + "spend_log": _spend(budget_key), + "key": key_spend if key_spend is not None else math.nan, + } + assert {name: math.isclose(cost, reference_spend) for name, cost in billed_costs.items()} == dict.fromkeys( + billed_costs, True + ), (billed_costs, reference_spend) + followup: Final = _send(gateway, "chat", invalid_model, budget_key) + assert followup.status_code == 200, followup.text + assert (len(invalid_wire.drain()), len(plain_wire.drain())) == (2, 1) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_provider_header_outside_latin1_does_not_fail_the_call(gateway: Gateway, tmp_path: Path) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + utf8_bytes_as_wire_text: Final = "東京リージョン".encode().decode("latin-1") + headers: Final = {**_request_id_headers(request_id), "x-provider-region": utf8_bytes_as_wire_text} + with ( + wire_server(_peer(headers)) as wire, + owned_proxy(gateway, tmp_path, {"DISABLE_AIOHTTP_TRANSPORT": "True"}) as httpx_transport, + httpx_transport.scenario() as scenario, + ): + models: Final = { + member: _MEMBERS[member](httpx_transport, scenario, wire) for member in ("bedrock_mantle", "groq") + } + key: Final = scenario.key(models=list(models.values())) + replies: Final = { + f"{endpoint}-{'stream' if stream else 'whole'}": _send( + httpx_transport, endpoint, models[member], key, stream=stream + ) + for endpoint, member, stream in _LATIN1_CALLS + } + assert {call: reply.status_code for call, reply in replies.items()} == dict.fromkeys(replies, 200), { + call: reply.text for call, reply in replies.items() + } + assert {call: reply.headers.get("llm_provider-x-provider-region") for call, reply in replies.items()} == ( + dict.fromkeys(replies, None) + ) + assert {call: _REPLY_TEXT in reply.text for call, reply in replies.items()} == dict.fromkeys(replies, True) + assert {endpoint: _text(endpoint, replies[f"{endpoint}-whole"]) for endpoint in _PATHS} == dict.fromkeys( + _PATHS, _REPLY_TEXT + ) + forwarding: Final = ("chat-whole", "chat-stream", "responses-whole") + assert {call: _llm_provider_request_ids(replies[call]) for call in forwarding} == dict.fromkeys( + forwarding, {"llm_provider-x-request-id": request_id, "llm_provider-x-amzn-requestid": request_id} + ) + assert len(wire.drain()) == len(_LATIN1_CALLS) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_an_unparseable_rate_limit_header_keeps_the_deployment_metrics(gateway: Gateway, tmp_path: Path) -> None: + multiproc_dir: Final = tmp_path / "prometheus" + multiproc_dir.mkdir() + with ( + wire_server(_peer(_RATE_LIMIT_HEADERS)) as wire, + owned_proxy( + gateway, + tmp_path, + {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)}, + config=series_cap_config(tmp_path, {}), + workers=2, + ) as metered, + metered.scenario() as scenario, + ): + calls: Final = { + f"{member}-{'stream' if stream else 'whole'}": (_MEMBERS[member](metered, scenario, wire), stream) + for member, stream in _METERED_CALLS + } + models: Final = {call: model for call, (model, _) in calls.items()} + key: Final = scenario.key(models=list(models.values())) + statuses: Final = { + call: _send(metered, "chat", model, key, stream=stream).status_code + for call, (model, stream) in calls.items() + } + assert statuses == dict.fromkeys(calls, 200) + observed: Final = eventually( + lambda: _deployment_metrics(scrape(metered), models), + lambda metrics: all(successes >= 1 for successes, _, _ in metrics.values()), + 30, + return_last_on_timeout=True, + ) + assert observed == dict.fromkeys(calls, (1.0, (0.0,), ())), observed + assert len(wire.drain()) == len(calls) diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index f12254b8d1d..ba2261e5a1f 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -8,10 +8,10 @@ while a client frame alone never flips a voice session into one. Every row runs which answers the injected and rewritten updates the way the vendors do. A blocked transcript on a transcription session reaches the client as a ``guardrail_violation`` error and -nothing else: the ``response.cancel``, the voiced block prompt and the ``response.create`` a voice session gets -never go to a backend that cannot speak, while ``on_violation: end_session`` and ``end_session_after_n_fails`` -still close the session. A guardrail that raises anything but a block closes the session the way it does on a -voice session. +nothing else: the ``response.cancel``, the blocked item's deletion, the voiced block prompt and the +``response.create`` a voice session gets never go to a backend that cannot speak, while +``on_violation: end_session`` and ``end_session_after_n_fails`` still close the session. A guardrail that raises +anything but a block closes the session the way it does on a voice session. """ from __future__ import annotations @@ -51,6 +51,7 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import ( OwnedProxy, UpstreamSlot, + graceful_stop_seconds, owned_proxy_process, owned_upstream, stop_root_process, @@ -60,7 +61,7 @@ from tests.integration.cost_calculation.cost_tracking_case import RealtimeRespon RecordProperty = Callable[[str, object], None] -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) WORKERS: Final = 2 BURST: Final = 20 @@ -83,7 +84,7 @@ CREATED_TYPES: Final = frozenset({"session.created", "transcription_session.crea TRANSCRIPT_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" RESPONSE_DONE: Final = "response.done" GUARDRAIL_VIOLATION: Final = ("guardrail_violation", "content_policy_violation") -MISSING_TURN_DETECTION_TYPE: Final = ("invalid_request_error", "missing_required_parameter") +SESSION_TYPE_MISMATCH: Final = ("invalid_request_error", "invalid_parameter") INVALID_SESSION_VALUE: Final = ("invalid_request_error", "invalid_value") FIVE_KB: Final = "x" * 5120 UNAUTHENTICATED_STATUS: Final = 403 @@ -104,10 +105,9 @@ BETA_TRANSCRIPTION_UPDATE: Final[dict[str, JsonValue]] = { "input_audio_transcription": {"model": TRANSCRIBE_MODEL}, } VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } RE_ENABLE_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", @@ -123,14 +123,32 @@ CLIENT_DECLARED_TRANSCRIPTION: Final[dict[str, JsonValue]] = { } CLIENT_DECLARED_TRANSCRIPTION_FORWARDED: Final[dict[str, JsonValue]] = { "type": "transcription", - "audio": {"input": {"turn_detection": {"create_response": False}}}, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." +REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." GA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } BETA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = {"turn_detection": {"type": "server_vad", "create_response": False}} +XAI_VOICE_MODEL: Final = "grok-voice" +GEMINI_LIVE_MODEL: Final = "gemini-3.8-live" +VERTEX_LIVE_LOCATION: Final = "us-central1" +UNTYPED_GATE: Final[dict[str, JsonValue]] = {"create_response": False} +SEMANTIC_VAD: Final[dict[str, JsonValue]] = {"type": "semantic_vad", "eagerness": "low"} +TUNED_SERVER_VAD: Final[dict[str, JsonValue]] = { + "type": "server_vad", + "prefix_padding_ms": 300, + "silence_duration_ms": 500, +} +GEMINI_TUNED_DETECTION: Final[dict[str, JsonValue]] = { + "disabled": True, + "prefixPaddingMs": 300, + "silenceDurationMs": 500, +} +CLIENT_CANCEL_ID: Final = "client_cancel_1" +CANCEL_NOT_ACTIVE: Final = ("invalid_request_error", "response_cancel_not_active") +AUTO_ANSWER_PREFIX: Final = "resp_auto_" COMMIT: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.commit"} APPEND: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.append", "audio": PCM_200MS_24KHZ} PUSH_TO_TALK: Final = "PUSH_TO_TALK" @@ -172,7 +190,13 @@ GUARDRAIL_END_CLOSE: Final = 1000 PROXY_FAILURE_CLOSE: Final = 1011 PROXY_FAILURE_REASON: Final = "proxy failed while relaying the upstream websocket" TOOL_OUTPUT_BLOCKED: Final = json.dumps({"error": "Tool output blocked by content policy"}) -VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") +VOICE_BLOCK_FRAMES: Final = ( + "response.cancel", + "conversation.item.delete", + "conversation.item.create", + "response.create", +) +XAI_VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") CREATED_SECONDS: Final = 60.0 STEP_SECONDS: Final = 15.0 SDK_SECONDS: Final = 45.0 @@ -318,6 +342,30 @@ def _voice_scenario(*transcripts: JsonValue) -> RealtimeResponse: return RealtimeResponse(content_type="application/x-realtime", events=tuple(_voice_turns(transcripts))) +def _xai_voice_scenario(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), dialect="xai" + ) + + +def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), created_after_first_frame=True + ) + + +def _live_scenario() -> RealtimeResponse: + return RealtimeResponse(content_type="application/x-realtime", events=()) + + +def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]: + return {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + + +def _without_auto_response(turn_detection: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {**turn_detection, "create_response": False} + + def _muse_scenario(transcript: str) -> RealtimeResponse: return RealtimeResponse( content_type="application/x-realtime", @@ -358,8 +406,41 @@ def _with_transcription_model(update: dict[str, JsonValue], model: str) -> dict[ return {**update, "audio": {"input": {**audio_input, "transcription": {"model": model}}}} -def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str) -> str: - return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base) +def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str, *, drop_params: bool) -> str: + return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base, drop_params=drop_params) + + +def _gemini_studio_deployment(scenario: Scenario, scenario_id: str, upstream_url: str) -> str: + return scenario.model( + model=f"gemini/{GEMINI_LIVE_MODEL}", + api_key=scenario_id, + api_base=upstream_url.rstrip("/"), + model_info={"mode": "realtime"}, + ) + + +def _vertex_credentials(upstream_url: str) -> str: + return json.dumps( + { + "type": "external_account", + "audience": "synthetic-vertex-audience", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": f"{upstream_url}/_oauth/token", + "credential_source": {"url": f"{upstream_url}/health"}, + } + ) + + +def _vertex_live_deployment(scenario: Scenario, project: str, upstream_url: str) -> str: + upstream: Final = upstream_url.rstrip("/") + return scenario.model( + model=f"vertex_ai/{GEMINI_LIVE_MODEL}", + vertex_project=project, + vertex_credentials=_vertex_credentials(upstream), + vertex_location=VERTEX_LIVE_LOCATION, + api_base=upstream, + model_info={"mode": "realtime"}, + ) def _named_deployment_id(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str: @@ -549,6 +630,55 @@ def _transcribe( return asyncio.run(_session(ws_base, path, query, key, frames, until=until, headers=headers)) +async def _received_after_sending( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...], until: str +) -> AsyncIterator[dict[str, JsonValue]]: + try: + for frame in frames: + await socket.send(json.dumps(frame)) + async for event in _until(socket, until, STEP_SECONDS): + yield event + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +def _is_proxy_answer(event: dict[str, JsonValue]) -> bool: + if event.get("type") != RESPONSE_DONE: + return False + return not string_value(object_value(event["response"])["id"]).startswith(AUTO_ANSWER_PREFIX) + + +async def _until_the_proxy_answer( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...] +) -> AsyncIterator[dict[str, JsonValue]]: + try: + yield JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), CREATED_SECONDS)) + for frame in frames: + await socket.send(json.dumps(frame)) + deadline: Final = asyncio.get_running_loop().time() + STEP_SECONDS + while (event := await _next_event(socket, deadline)) is not None: + yield event + if _is_proxy_answer(event): + return + yield {"type": "timeout"} + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +async def _talked_until_the_proxy_answer( + ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...] +) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _until_the_proxy_answer(socket, frames)]), None) + + +async def _talked_at_connect(ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...]) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _received_after_sending(socket, frames, RESPONSE_DONE)]), None) + + def _talk( ws_base: str, query: str, @@ -617,10 +747,43 @@ def _session_updates(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonVa return tuple(frame.get("session") for frame in _sent(observed) if frame.get("type") == "session.update") +def _assert_voice_updates(observed: tuple[dict[str, JsonValue], ...], *later: dict[str, JsonValue]) -> None: + updates: Final = _session_updates(observed) + assert updates[0] == GA_INJECTED_UPDATE, updates + assert updates[1] == VOICE_UPDATE_GATED, updates + assert updates[2:] == later, updates + + def _query(upgrade: dict[str, JsonValue]) -> JsonValue: return upgrade["query"] +def _declared_turn_detection(update: JsonValue) -> JsonValue: + session: Final = object_value(update) + audio_input: Final = object_value(object_value(session.get("audio", {})).get("input", {})) + return audio_input["turn_detection"] if "turn_detection" in audio_input else session.get("turn_detection") + + +def _upstream_turn_detections(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonValue, ...]: + return tuple(_declared_turn_detection(update) for update in _session_updates(observed)) + + +def _response_ids(session: Session) -> tuple[str, ...]: + done: Final = tuple(object_value(event["response"]) for event in session.events if event["type"] == RESPONSE_DONE) + return tuple(string_value(response["id"]) for response in done) + + +def _error_event_ids(session: Session) -> tuple[JsonValue, ...]: + errors: Final = tuple(object_value(event["error"]) for event in session.events if event["type"] == "error") + return tuple(error.get("event_id") for error in errors) + + +def _setup_detection(observed: tuple[dict[str, JsonValue], ...]) -> JsonValue: + setups: Final = tuple(object_value(frame["setup"]) for frame in _sent(observed) if "setup" in frame) + assert len(setups) == 1, _sent(observed) + return object_value(setups[0].get("realtimeInputConfig", {})).get("automaticActivityDetection") + + def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)), @@ -704,6 +867,7 @@ def _write_config(directory: Path, ca_bundle: Path) -> Path: "litellm_settings": { "ssl_verify": str(ca_bundle), "enable_redis_auth_cache": True, + "gemini_live_defer_setup": True, "cache": True, "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, }, @@ -918,19 +1082,18 @@ def test_voice_session_keeps_the_injected_update_and_the_clean_transcript_flows( assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.errors == (), session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", "response.create", ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) rows: Final = _spend_rows(key, 1) assert rows[0]["call_type"] == "_arealtime", rows @@ -945,17 +1108,15 @@ def test_voice_session_blocked_transcript_is_refused_through_the_backend(guardra ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION,), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardrail_proxy: OwnedProxy) -> None: @@ -967,12 +1128,8 @@ def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardr session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.types[-1] == RESPONSE_DONE, session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert _session_updates(observed) == ( - GA_INJECTED_UPDATE, - VOICE_UPDATE_FORWARDED, - RE_ENABLE_FORWARDED, - ), _session_updates(observed) + assert session.errors == (), session + _assert_voice_updates(observed, RE_ENABLE_FORWARDED) def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail(guardrail_proxy: OwnedProxy) -> None: @@ -984,16 +1141,14 @@ def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail( session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames, seconds=20) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[0] == TURN_DETECTION_TYPE_MISSING, session + assert session.errors == (SESSION_TYPE_MISMATCH, GUARDRAIL_VIOLATION), session + assert session.error_messages[0] == REALTIME_UPDATE_REFUSED, session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) assert _session_updates(observed) == ( GA_INJECTED_UPDATE, @@ -1014,6 +1169,202 @@ def test_backend_session_created_typed_transcription_skips_the_injection_on_the_ assert [_query(upgrade) for upgrade in _upgrades(observed)] == [[["model", TRANSCRIBE_MODEL]]] +def test_voice_session_typed_turn_detection_keeps_its_type_and_tuning(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(SEMANTIC_VAD)) + session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(SEMANTIC_VAD)), + ), _session_updates(observed) + + +@dataclass(frozen=True, slots=True) +class LaterUpdateCase: + headers: Mapping[str, str] + first: dict[str, JsonValue] + later: dict[str, JsonValue] + forwarded: tuple[dict[str, JsonValue], ...] + + +LATER_UPDATE_CASES: Final = ( + pytest.param( + LaterUpdateCase( + {}, + VOICE_UPDATE, + {"type": "realtime", "instructions": "later"}, + (GA_INJECTED_UPDATE, VOICE_UPDATE_GATED), + ), + id="ga", + ), + pytest.param( + LaterUpdateCase( + BETA_HEADERS, + {"turn_detection": {"type": "server_vad"}}, + {"instructions": "later"}, + (BETA_INJECTED_UPDATE, BETA_INJECTED_UPDATE), + ), + id="beta", + ), +) + + +@pytest.mark.parametrize("case", LATER_UPDATE_CASES) +def test_later_update_without_turn_detection_is_forwarded_verbatim( + guardrail_proxy: OwnedProxy, case: LaterUpdateCase +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(case.first), _update_frame(case.later), COMMIT, until=RESPONSE_DONE),), + headers=case.headers, + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == (*case.forwarded, case.later), _session_updates(observed) + + +def test_session_update_sent_before_session_created_follows_the_gate(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario_created_after_the_first_frame(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_at_connect(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.types == ( + "session.created", + "session.updated", + "session.updated", + TRANSCRIPT_COMPLETED, + RESPONSE_DONE, + ), session + assert session.errors == (), session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(TUNED_SERVER_VAD)), + ), _session_updates(observed) + + +def test_client_response_cancel_error_still_reaches_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + cancel: Final[dict[str, JsonValue]] = {"type": "response.cancel", "event_id": CLIENT_CANCEL_ID} + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, (_step(cancel, until="error"),) + ) + assert session.errors == (CANCEL_NOT_ACTIVE,), session + assert _error_event_ids(session) == (CLIENT_CANCEL_ID,), session + + +def _xai_blocked_session( + guardrail_proxy: OwnedProxy, scenario: Scenario +) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: + handle: Final = _scripted(scenario, _xai_voice_scenario(BLOCKED_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = scenario.model(model=f"xai/{XAI_VOICE_MODEL}", api_key=handle.scenario_id) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_until_the_proxy_answer(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + return session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + + +def test_xai_voice_session_gets_untyped_gates_and_no_unguarded_answer(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + turn_detections: Final = _upstream_turn_detections(observed) + assert turn_detections == (UNTYPED_GATE, UNTYPED_GATE), turn_detections + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert GUARDRAIL_VIOLATION in session.errors, session + assert not any(response.startswith(AUTO_ANSWER_PREFIX) for response in _response_ids(session)), session + + +def test_xai_voice_session_block_sends_xai_no_item_delete(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert _sent_types(observed) == ( + "session.update", + "session.update", + "input_audio_buffer.commit", + *XAI_VOICE_BLOCK_FRAMES, + ), _sent(observed) + + +@pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason=( + "xAI answers the guardrail's response.cancel with an error carrying its own event_id, so the proxy cannot " + "tell it apart from an error for a client event and forwards it, as the merge base does" + ), +) +def test_xai_guardrail_cancel_error_does_not_reach_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, _ = _xai_blocked_session(guardrail_proxy, scenario) + assert session.errors == (GUARDRAIL_VIOLATION,), session + + +TUNED_VOICE_UPDATES: Final = ( + pytest.param(_realtime_update(TUNED_SERVER_VAD), id="nested"), + pytest.param({"turn_detection": TUNED_SERVER_VAD}, id="flat"), +) + + +@pytest.mark.parametrize("update", TUNED_VOICE_UPDATES) +def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guardrail( + guardrail_proxy: OwnedProxy, update: dict[str, JsonValue] +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _gemini_studio_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(update), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == {"disabled": True}, _sent(observed) + + +def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _vertex_live_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed) + + MUSE_PUSH_TO_TALK_FRAMES: Final[tuple[dict[str, JsonValue], ...]] = ( {"binary_bytes": MUSE_PACKET_BYTES}, {"binary_bytes": MUSE_PACKET_BYTES}, @@ -1041,10 +1392,11 @@ def _muse_session( turn_detection: JsonValue, *, until: str | None = None, + drop_params: bool = False, ) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: handle: Final = _scripted(scenario, _muse_scenario(transcript), control_url=tls_upstream) key: Final = scenario.key() - model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream) + model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream, drop_params=drop_params) verdict: Final = TRANSCRIPT_COMPLETED if BLOCKED_WORD not in transcript else "error" session: Final = _talk( _ws_base(_owned_url(guardrail_proxy)), @@ -1120,6 +1472,40 @@ def test_muse_server_vad_transcription_session_keeps_endpointing( ), _muse_audio_frames(observed) +def test_muse_semantic_vad_under_drop_params_falls_back_to_endpointing( + guardrail_proxy: OwnedProxy, tls_upstream: str +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + f"&{TRANSCRIPTION_QUERY}", + {"type": "semantic_vad"}, + drop_params=True, + ) + assert session.transcripts == (CLEAN_TRANSCRIPT,), session + assert session.errors == (), session + assert _muse_handshake(observed)["mode"] == ENDPOINTING, _muse_handshake(observed) + + +def test_muse_semantic_vad_without_drop_params_is_refused(guardrail_proxy: OwnedProxy, tls_upstream: str) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + "", + {"type": "semantic_vad"}, + until=TRANSCRIPT_COMPLETED, + ) + assert session.types[-1] == "closed", session + assert session.transcripts == (), session + assert _upgrades(observed) == (), observed + + @pytest.mark.xfail( strict=True, reason=( @@ -1229,8 +1615,8 @@ def test_opt_in_guardrail_gates_a_voice_session_on_an_opted_out_key(guardrail_pr ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + assert session.errors == (GUARDRAIL_VIOLATION,), session + _assert_voice_updates(observed) def test_opt_in_guardrail_leaves_a_transcription_session_alone_on_an_opted_out_key(guardrail_proxy: OwnedProxy) -> None: @@ -1717,12 +2103,11 @@ def test_voice_session_guardrail_raising_runtime_error_closes_with_the_same_prox assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", ), session - assert session.errors[0] == MISSING_TURN_DETECTION_TYPE, session _assert_closed_by_a_proxy_failure(session) assert _sent_types(observed) == ("session.update", "session.update", "input_audio_buffer.commit"), _sent( observed @@ -1740,8 +2125,8 @@ def test_voice_session_guardrail_raising_value_error_is_voiced_through_the_backe (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE),), ) assert session.transcripts == (blocked,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[1] == f"{VALUE_ERROR_WORD} is not allowed", session + assert session.errors == (GUARDRAIL_VIOLATION,), session + assert session.error_messages == (f"{VALUE_ERROR_WORD} is not allowed",), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", @@ -1764,7 +2149,7 @@ def test_voice_session_second_violation_under_end_session_after_n_fails_closes_a (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE), _step(COMMIT, until="closed")), ) assert session.transcripts == (blocked, blocked), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session assert session.types[-1] == "closed", session assert session.close_code == GUARDRAIL_END_CLOSE, session assert _sent_types(observed) == ( @@ -1948,7 +2333,7 @@ def test_non_string_transcript_field_closes_a_voice_session_with_the_same_proxy_ assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index 4dd7fd9aeef..c7489685273 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -114,12 +114,10 @@ TRANSCRIPT_FIELDS: Final = ( pytest.param({"transcript": {"text": "a"}}, id="object"), ) VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, @@ -982,15 +980,15 @@ def test_list_mode_transcript_guardrail_still_turns_auto_response_off_and_answer assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.error_messages == (), session assert _sent(observed) == ( _update_frame(INJECTED_UPDATE), - _update_frame(VOICE_UPDATE_FORWARDED), + _update_frame(VOICE_UPDATE_GATED), COMMIT, PROXY_RESPONSE_CREATE, ), _sent(observed) 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/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 8de93256f6c..ec3fb941037 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -4,18 +4,22 @@ import threading from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias, get_args from uuid import uuid4 import anthropic import httpx import pytest -from integration._support.client import Gateway, Scenario -from integration._support.process import graceful_stop_seconds +from integration._support.client import Gateway, Scenario, object_value, string_value +from integration._support.process import 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 +from typing_extensions import assert_never pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -302,6 +306,396 @@ def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provide assert peer.drain() == () +_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180 +_FILES: Final = 3 +_READ_FILE_TOOL: Final[JsonValue] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_OLDER_TOOL_RESULTS: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"{_FILE_TEXT}", + } + ], + }, + ) + + +def _agent_body(model: str, marker: str, options: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return { + "model": model, + "tools": [_READ_FILE_TOOL], + "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], + "context_management": _CLEAR_OLDER_TOOL_RESULTS, + **options, + } + + +def _files_kept(contents: JsonValue, marker: str) -> tuple[bool, ...]: + sent: Final = json.dumps(contents) + return tuple(f"" in sent for index in range(_FILES)) + + +def _contents_sent_upstream( + gateway: Gateway, peer: _GeminiPeer, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[JsonValue, JsonValue]: + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=headers) + assert generated.status_code == 200, generated.text + generate_requests: Final = tuple(request for request in peer.drain() if request.target.endswith(":generateContent")) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=headers) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + count_requests: Final = peer.drain() + assert len(generate_requests) == 1 + assert [request.target.endswith(":countTokens") for request in count_requests] == [True] + return _json_body(generate_requests[0])["contents"], _counted_contents(_json_body(count_requests[0])) + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_counts_the_history_messages_sends_after_context_edits(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, {})) + assert _files_kept(generated, marker) == (False, False, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "files_kept"), + ( + ({"additional_drop_params": ["context_management"]}, {}, (True, True, True)), + ({}, {"additional_drop_params": ["context_management"]}, (True, True, True)), + ({"additional_drop_params": ["context_management"]}, {"additional_drop_params": []}, (False, False, True)), + ), + ids=("deployment-drops", "request-drops", "request-overrides-the-deployment-drop"), +) +def test_count_tokens_follows_the_context_management_opt_out_messages_follows( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + files_kept: tuple[bool, ...], +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, request_options)) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +_RoutedCase: TypeAlias = Literal[ + "router-default-drops", + "fallback-deployment-drops", + "fallback-router-default-drops", + "fallback-request-drops", + "alias-to-a-litellm-model-drops", + "completion-model-drops", +] +_DROPS_CONTEXT_MANAGEMENT: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"additional_drop_params": ["context_management"]} +) + + +@dataclass(frozen=True, slots=True) +class _RoutedProxy: + settings: Mapping[str, JsonValue] + requested: str + request_options: Mapping[str, JsonValue] + + +def _scripted_gemini(model_name: str, model: str, api_base: str, options: Mapping[str, JsonValue]) -> JsonValue: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "scripted-gemini-key", **options}, + } + + +def _default_fallback_settings( + prefix: str, api_base: str, deployment_options: Mapping[str, JsonValue], router_defaults: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + return { + "model_list": [ + _scripted_gemini(f"{prefix}-fallback", f"gemini/{_MODEL}", api_base, deployment_options), + _scripted_gemini(f"{prefix}-wild/*", "gemini/*", api_base, {}), + ], + "router_settings": { + "default_fallbacks": [f"{prefix}-fallback"], + **({"default_litellm_params": dict(router_defaults)} if router_defaults else {}), + }, + } + + +def _routed_proxy(case: _RoutedCase, prefix: str, api_base: str) -> _RoutedProxy: + match case: + case "router-default-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, {})], + "router_settings": {"default_litellm_params": dict(_DROPS_CONTEXT_MANAGEMENT)}, + }, + prefix, + {}, + ) + case "fallback-deployment-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, _DROPS_CONTEXT_MANAGEMENT, {}), f"{prefix}-unrouted", {} + ) + case "fallback-router-default-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, _DROPS_CONTEXT_MANAGEMENT), f"{prefix}-unrouted", {} + ) + case "fallback-request-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, {}), f"{prefix}-unrouted", _DROPS_CONTEXT_MANAGEMENT + ) + case "alias-to-a-litellm-model-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT)], + "router_settings": {"model_group_alias": {f"{prefix}-alias": f"gemini/{_MODEL}"}}, + }, + f"{prefix}-alias", + {}, + ) + case "completion-model-drops": + return _RoutedProxy( + { + "model_list": [ + _scripted_gemini(f"{prefix}-keeps", f"gemini/{_MODEL}", api_base, {}), + _scripted_gemini(f"{prefix}-drops", f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT), + ], + "general_settings": {"completion_model": f"{prefix}-drops"}, + }, + f"{prefix}-keeps", + {}, + ) + case _: + assert_never(case) + + +@pytest.mark.parametrize("case", get_args(_RoutedCase)) +def test_count_tokens_follows_the_opt_out_of_the_deployment_proxy_settings_send_to( + gateway: Gateway, tmp_path: Path, case: _RoutedCase +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire: + routed: Final = _routed_proxy(case, f"routed-{uuid4().hex}", wire.url) + config: Final = tmp_path / "routed.yaml" + config.write_text(json.dumps(routed.settings), encoding="utf-8") + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + owned.gateway, peer, _agent_body(routed.requested, marker, routed.request_options) + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def _named_deployment( + scenario: Scenario, + model_name: str, + litellm_params: Mapping[str, JsonValue], + model_info: Mapping[str, JsonValue], +) -> str: + created: Final = scenario.gateway.post( + "/model/new", + {"model_name": model_name, "litellm_params": dict(litellm_params), "model_info": dict(model_info)}, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def test_count_tokens_reads_the_opt_out_of_a_deployment_id_beside_a_wildcard_route(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"integration-{uuid4().hex}" + _named_deployment( + scenario, + f"{prefix}-*", + {"model": "gemini/*", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + deployment_id: Final = _named_deployment( + scenario, + prefix, + { + "model": f"gemini/{_MODEL}", + "api_base": wire.url, + "api_key": "scripted-gemini-key", + "additional_drop_params": ["context_management"], + }, + {"id": f"{prefix}-dropping"}, + ) + assert deployment_id == f"{prefix}-dropping" + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(deployment_id, marker, {})) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("requested", "wildcards", "files_kept"), + ( + ("dropping", (), (True, True, True)), + ("editing", (), (False, False, True)), + ("dropping", ("*",), (False, False, True)), + ), + ids=("deployment-drops", "sibling-deployment-drops", "wildcard-route-matches-first"), +) +def test_count_tokens_reads_the_opt_out_of_the_deployment_its_litellm_model_name_routes_to( + gateway: Gateway, requested: str, wildcards: tuple[str, ...], files_kept: tuple[bool, ...] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"gemini/{_MODEL}-{uuid4().hex}" + scenario.model( + model=f"{prefix}-dropping", + api_base=wire.url, + api_key="scripted-gemini-key", + additional_drop_params=["context_management"], + ) + scenario.model(model=f"{prefix}-editing", api_base=wire.url, api_key="scripted-gemini-key") + for wildcard in wildcards: + _named_deployment( + scenario, + f"{prefix}-{wildcard}", + {"model": f"{prefix}-{wildcard}", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(f"{prefix}-{requested}", marker, {})) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +def _error_type(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["type"] + + +def _error_message(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["message"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_a_context_edit_messages_rejects_without_an_upstream_call( + gateway: Gateway, provider: str +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + body: Final = _agent_body(model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +def test_count_tokens_rejects_an_invalid_edit_from_a_coding_agent_that_turns_drop_params_off( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + body: Final = _agent_body( + model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, "drop_params": False} + ) + claude_code: Final = MappingProxyType({"User-Agent": "claude-cli/2.1.0 (external, cli)"}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=claude_code) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=claude_code) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "user_agent"), + ( + ({"drop_params": True}, {}, "python-httpx"), + ({}, {"drop_params": True}, "python-httpx"), + ({}, {}, "claude-cli/2.1.0 (external, cli)"), + ({}, {}, "codex_cli_rs/0.50.0"), + ), + ids=("deployment-drop-params", "request-drop-params", "claude-code-client", "codex-client"), +) +def test_count_tokens_counts_the_unedited_body_messages_sends_when_drop_params_drops_an_invalid_edit( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + user_agent: str, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + gateway, + peer, + _agent_body(model, marker, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, **request_options}), + {"User-Agent": user_agent}, + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def test_count_tokens_answers_for_the_unedited_body_when_a_context_edit_fails_inside_the_proxy( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + unedited: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + edited: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _UNCONVERTIBLE_MESSAGES, "context_management": _CLEAR_OLDER_TOOL_RESULTS}, + ) + assert unedited.status_code == 400, unedited.text + assert (edited.status_code, _error_type(edited), _error_message(edited)) == ( + unedited.status_code, + _error_type(unedited), + _error_message(unedited), + ) + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: 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/sdk/conftest.py b/tests/integration/sdk/conftest.py index 066cca6250a..05f4f3d7d3a 100644 --- a/tests/integration/sdk/conftest.py +++ b/tests/integration/sdk/conftest.py @@ -1,7 +1,11 @@ +from collections.abc import AsyncIterator from typing import Final -import litellm import pytest +import pytest_asyncio + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER CALLBACK_LISTS: Final = ( "callbacks", @@ -14,7 +18,12 @@ CALLBACK_LISTS: Final = ( ) -@pytest.fixture(autouse=True) -def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest_asyncio.fixture(autouse=True, loop_scope="function") +async def isolate_litellm_callback_lists(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[None]: for name in CALLBACK_LISTS: monkeypatch.setattr(litellm, name, list(getattr(litellm, name))) + try: + yield + finally: + await GLOBAL_LOGGING_WORKER.flush() + await GLOBAL_LOGGING_WORKER.stop() diff --git a/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py b/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py new file mode 100644 index 00000000000..99f5fa941f7 --- /dev/null +++ b/tests/integration/sdk/test_responses_bridge_sync_reasoning_sdk.py @@ -0,0 +1,130 @@ +import json +import uuid +from collections.abc import Iterable, Mapping, Sequence +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +import litellm +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator + +_BACKEND: Final = "qwen3-sync-bridge-audit" +_API_KEY: Final = "synthetic-hosted-vllm-key" +_ANSWER: Final = "It is sunny in Paris." +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_ITEM_EVENTS: Final = frozenset({"response.output_item.added", "response.output_item.done"}) + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Item(_Frozen): + type: str + + +class _Event(_Frozen): + type: str + item: _Item | None = None + output_index: int | None = None + delta: JsonValue = None + + +def _frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None = None) -> bytes: + usage: Final = ( + {"usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}} if finish_reason else {} + ) + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + **usage, + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _sse(deltas: Sequence[Mapping[str, JsonValue]]) -> Reply: + identity: Final = f"chatcmpl-sync-bridge-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + _frame(identity, {"role": "assistant", "content": ""}), + *(_frame(identity, delta) for delta in deltas), + _frame(identity, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + + +def _streamed_chat_posts(wire: Wire) -> list[JsonValue]: + posted: Final = [request for request in wire.drain() if request.method == "POST"] + assert [request.target for request in posted] == ["/v1/chat/completions"] + return [_JSON_OBJECT.validate_json(request.body).get("stream") for request in posted] + + +def _events(stream: Iterable[object]) -> tuple[_Event, ...]: + return tuple(_Event.model_validate(_JSON_OBJECT.validate_python(_dumped(event))) for event in stream) + + +def _dumped(event: object) -> object: + assert isinstance(event, BaseModel), type(event) + return event.model_dump(mode="json") + + +def _delta_text(events: Sequence[_Event], event_type: str) -> str: + return "".join(str(event.delta) for event in events if event.type == event_type) + + +@pytest.mark.parametrize( + "deltas", + [ + pytest.param( + ({"reasoning_content": "REASONING"}, {"content": _ANSWER}), + id="reasoning-then-answer-chunks", + ), + pytest.param(({"reasoning_content": "REASONING", "content": _ANSWER},), id="reasoning-and-answer-in-one-chunk"), + ], +) +def test_sync_responses_stream_closes_the_reasoning_item_before_the_message_opens( + deltas: tuple[dict[str, JsonValue], ...], +) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Sync bridge reasoning {marker}." + scripted: Final = tuple( + {key: reasoning if key == "reasoning_content" else value for key, value in delta.items()} for delta in deltas + ) + + def respond(request: Request) -> Reply: + return Reply(body=b'{"object":"list","data":[]}') if request.method == "GET" else _sse(scripted) + + with wire_server(respond) as wire: + stream: Final = litellm.responses( + model=f"hosted_vllm/{_BACKEND}", + input=f"Weather in Paris? {marker}", + stream=True, + api_base=wire.url + "/v1", + api_key=_API_KEY, + use_chat_completions_api=True, + num_retries=0, + ) + assert isinstance(stream, LiteLLMCompletionStreamingIterator), type(stream) + events: Final = _events(stream) + + assert _streamed_chat_posts(wire) == [True] + lifecycle: Final = [ + (event.type, event.item.type, event.output_index) + for event in events + if event.type in _ITEM_EVENTS and event.item is not None + ] + assert lifecycle == [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", "message", 1), + ("response.output_item.done", "message", 1), + ], [event.type for event in events] + assert _delta_text(events, "response.reasoning_summary_text.delta") == reasoning + assert _delta_text(events, "response.output_text.delta") == _ANSWER diff --git a/tests/integration/sdk/test_router_budget_limiter_wire.py b/tests/integration/sdk/test_router_budget_limiter_wire.py index cf4fda41891..0863ee6d6cc 100644 --- a/tests/integration/sdk/test_router_budget_limiter_wire.py +++ b/tests/integration/sdk/test_router_budget_limiter_wire.py @@ -1,7 +1,9 @@ from __future__ import annotations +import asyncio import json import os +import time import uuid from collections.abc import Iterator from typing import Final @@ -222,3 +224,95 @@ async def test_tag_budget_blocks_exhausted_tag_but_allows_another( assert response.choices[0].message.content == "openai" assert len(requests) == 1 assert json.loads(requests[0].body)["model"] == _MODEL + + + +_ISOLATED_SPEND_KEY: Final = "provider_spend:hosted_vllm:1d" +_INPUT_COST_PER_TOKEN: Final = 0.001 +_OUTPUT_COST_PER_TOKEN: Final = 0.002 + + +def _isolated_provider_deployment(api_base: str) -> dict[str, JsonValue]: + """Only these two tests route to hosted_vllm, so success callbacks still draining from earlier tests cannot land here.""" + return { + "model_name": "isolated-budget", + "litellm_params": { + "model": "hosted_vllm/llama-3", + "api_base": api_base, + "api_key": _API_KEY, + "input_cost_per_token": _INPUT_COST_PER_TOKEN, + "output_cost_per_token": _OUTPUT_COST_PER_TOKEN, + }, + "model_info": {"id": "vllm-deployment"}, + } + + +def _isolated_budget_router(api_base: str, redis_client: Redis | None) -> Router: + redis_params: Final = ( + {} + if redis_client is None + else {"redis_host": os.environ["REDIS_HOST"], "redis_port": int(os.environ["REDIS_PORT"])} + ) + return Router( + model_list=[_isolated_provider_deployment(api_base)], + provider_budget_config={"hosted_vllm": BudgetConfig(budget_duration="1d", max_budget=10.0)}, + num_retries=0, + disable_cooldowns=True, + **redis_params, + ) + + +async def _complete_and_wait_for_spend(router: Router, previous: float) -> float: + response: Final = await router.acompletion( + model="isolated-budget", + messages=[{"role": "user", "content": "budget tracking"}], + ) + cost: Final = ( + response.usage.prompt_tokens * _INPUT_COST_PER_TOKEN + response.usage.completion_tokens * _OUTPUT_COST_PER_TOKEN + ) + deadline: Final = time.monotonic() + 5.0 + while time.monotonic() < deadline: + spend: float | None = router.cache.in_memory_cache.get_cache(_ISOLATED_SPEND_KEY) + if spend is not None and spend > previous: + assert spend == pytest.approx(previous + cost) + return spend + await asyncio.sleep(0.05) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} never rose above {previous}") + + +@pytest.mark.asyncio +async def test_router_without_redis_tracks_spend_but_queues_no_redis_increments() -> None: + """A Router with no Redis never flushes, so anything it queues for Redis is retained for the life of the process.""" + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=None) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + third: Final = await _complete_and_wait_for_spend(router, previous=second) + wire.drain() + + assert third > second > first > 0.0 + assert router.router_budget_logger is not None + assert router.router_budget_logger.redis_increment_operation_queue == [] + + +async def _wait_for_redis_spend(redis_client: Redis, at_least: float) -> float: + deadline: Final = time.monotonic() + 15.0 + while time.monotonic() < deadline: + raw: str | None = redis_client.get(_ISOLATED_SPEND_KEY) + if raw is not None and float(raw) >= at_least: + return float(raw) + await asyncio.sleep(0.1) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} in Redis never reached {at_least}") + + +@pytest.mark.asyncio +async def test_router_with_redis_flushes_queued_spend_increments(redis_client: Redis) -> None: + redis_client.delete(_ISOLATED_SPEND_KEY) + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=redis_client) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + redis_spend: Final = await _wait_for_redis_spend(redis_client, at_least=second) + wire.drain() + + assert redis_spend == pytest.approx(second) 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_fake_stream_reasoning_wire.py b/tests/integration/streaming/test_fake_stream_reasoning_wire.py new file mode 100644 index 00000000000..e46208365dd --- /dev/null +++ b/tests/integration/streaming/test_fake_stream_reasoning_wire.py @@ -0,0 +1,874 @@ +import base64 +import json +import math +import os +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from importlib.resources import files +from itertools import chain +from typing import Final +from urllib.parse import unquote + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.upstream import ( + _aws_event_frame, # pyright: ignore[reportPrivateUsage] # the scripted Bedrock peer's only event-stream encoder +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from redis import Redis + +_TOKEN: Final = "synthetic-bedrock-bearer" +_NATIVE_REDACTED: Final = base64.b64encode(b"\x12\x1esynthetic redacted reasoning payload").decode() +_CONVERSE_REDACTED: Final = base64.b64encode(_NATIVE_REDACTED.encode()).decode() +_SIGNATURE: Final = "EqoBCkgIBxABGAIiQsyntheticSignature" +_ANSWER: Final[dict[str, JsonValue]] = {"city": "Tokyo", "population_millions": 14.0} +_THINKING_TOKENS: Final = 87 +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_THINKING: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}, "population_millions": {"type": "number"}}, + "required": ["city", "population_millions"], + "additionalProperties": False, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "city_fact", "schema": _SCHEMA}, +} +_CONVERSE_USAGE: Final[dict[str, JsonValue]] = { + "inputTokens": 40, + "outputTokens": 120, + "cacheReadInputTokens": 50, + "totalTokens": 210, +} +_LLAMA_MODEL_ID: Final = "us.meta.llama3-3-70b-instruct-v1:0" +_VLLM_BACKEND: Final = "qwen3-fake-stream-audit" +_VLLM_KEY: Final = "synthetic-hosted-vllm-key" +_AZURE_DEPLOYMENT: Final = "reasoning-deployment" +_AZURE_API_VERSION: Final = "2025-03-01-preview" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_DISCOVERY_PROBE: Final = ("GET", "/v1/models") + + +def _tool(name: str) -> dict[str, JsonValue]: + return { + "type": "function", + "function": { + "name": name, + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + } + + +class _Frozen(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _CostEntry(_Frozen): + litellm_provider: str | None = None + supports_native_structured_output: bool | None = None + supports_reasoning: bool | None = None + + +class _ToolFunction(_Frozen): + name: str | None = None + arguments: str | None = None + + +class _ToolCall(_Frozen): + index: int = 0 + function: _ToolFunction + + +class _PromptDetails(_Frozen): + cached_tokens: int | None = None + + +class _CompletionDetails(_Frozen): + reasoning_tokens: int | None = None + + +class _Usage(_Frozen): + prompt_tokens: int + completion_tokens: int + total_tokens: int + prompt_tokens_details: _PromptDetails | None = None + completion_tokens_details: _CompletionDetails | None = None + + +class _Message(_Frozen): + content: str | None = None + reasoning_content: str | None = None + thinking_blocks: tuple[dict[str, JsonValue], ...] | None = None + tool_calls: tuple[_ToolCall, ...] | None = None + + +class _Choice(_Frozen): + message: _Message + finish_reason: str | None = None + + +class _Completion(_Frozen): + choices: tuple[_Choice, ...] + usage: _Usage + + +class _StreamChoice(_Frozen): + delta: _Message + finish_reason: str | None = None + + +class _Chunk(_Frozen): + choices: tuple[_StreamChoice, ...] = () + usage: _Usage | None = None + + +class _ConverseMessage(_Frozen): + role: str + content: tuple[dict[str, JsonValue], ...] + + +class _ConverseRequest(_Frozen): + messages: tuple[_ConverseMessage, ...] + + +@dataclass(frozen=True, slots=True) +class _Answer: + content: str + reasoning: str + thinking_blocks: tuple[dict[str, JsonValue], ...] + tool_calls: tuple[tuple[str, JsonValue], ...] + finish_reason: str | None + usage: _Usage | None + + +_COST_MAP: Final = TypeAdapter(dict[str, _CostEntry]) + + +def _fakes_structured_output_stream(model: str, entry: _CostEntry) -> bool: + return ( + model.startswith("us.anthropic.claude") + and entry.litellm_provider == "bedrock_converse" + and entry.supports_native_structured_output is False + and entry.supports_reasoning is True + ) + + +def _claude_without_native_structured_output() -> str: + entries: Final = _COST_MAP.validate_json( + (files("litellm") / "model_prices_and_context_window_backup.json").read_bytes() + ) + return next(model for model, entry in entries.items() if _fakes_structured_output_stream(model, entry)) + + +def _converse_reply(content: Sequence[JsonValue], stop_reason: str) -> Reply: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": list(content)}}, + "stopReason": stop_reason, + "usage": _CONVERSE_USAGE, + "additionalModelResponseFields": { + "usage": {"output_tokens_details": {"thinking_tokens": _THINKING_TOKENS}} + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _reasoned_json_answer(reasoning: str) -> list[JsonValue]: + return [ + {"reasoningContent": {"reasoningText": {"text": reasoning, "signature": _SIGNATURE}}}, + {"reasoningContent": {"redactedContent": _CONVERSE_REDACTED}}, + {"toolUse": {"toolUseId": "tooluse_json_1", "name": "json_tool_call", "input": _ANSWER}}, + ] + + +def _is_replay(request: Request) -> bool: + return any(message.role == "assistant" for message in _ConverseRequest.model_validate_json(request.body).messages) + + +def _converse_peer(reasoning: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse"), request.target + if _is_replay(request): + return _converse_reply([{"text": "replayed"}], "end_turn") + return _converse_reply(_reasoned_json_answer(reasoning), "tool_use") + + return respond + + +def _converse_targets(wire: Wire) -> list[str]: + return [unquote(request.target).rsplit("/", 1)[-1] for request in wire.drain()] + + +def _claude_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{_claude_without_native_structured_output()}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _parsed_arguments(arguments: str | None) -> JsonValue: + return _JSON_OBJECT.validate_json(arguments or "{}") + + +def _completion_answer(completion: _Completion) -> _Answer: + choice: Final = completion.choices[0] + return _Answer( + content=choice.message.content or "", + reasoning=choice.message.reasoning_content or "", + thinking_blocks=choice.message.thinking_blocks or (), + tool_calls=tuple( + (call.function.name or "", _parsed_arguments(call.function.arguments)) + for call in choice.message.tool_calls or () + ), + finish_reason=choice.finish_reason, + usage=completion.usage, + ) + + +def _joined_call(calls: Sequence[_ToolCall], index: int) -> tuple[str, JsonValue]: + pieces: Final = tuple(call.function for call in calls if call.index == index) + return ( + "".join(piece.name or "" for piece in pieces), + _parsed_arguments("".join(piece.arguments or "" for piece in pieces)), + ) + + +def _streamed_answer(chunks: Sequence[_Chunk]) -> _Answer: + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + deltas: Final = tuple(choice.delta for choice in choices) + calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas)) + finishes: Final = tuple(choice.finish_reason for choice in choices if choice.finish_reason) + usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) + return _Answer( + content="".join(delta.content or "" for delta in deltas), + reasoning="".join(delta.reasoning_content or "" for delta in deltas), + thinking_blocks=tuple(chain.from_iterable(delta.thinking_blocks or () for delta in deltas)), + tool_calls=tuple(_joined_call(calls, index) for index in sorted({call.index for call in calls})), + finish_reason=finishes[-1] if finishes else None, + usage=usages[-1] if usages else None, + ) + + +def _chat(gateway: Gateway, body: Mapping[str, JsonValue]) -> _Answer: + completion: Final = _openai_client(gateway).chat.completions.create( + model=str(body["model"]), + messages=[{"role": "user", "content": str(body["question"])}], + extra_body={key: value for key, value in body.items() if key not in ("model", "question")}, + ) + return _completion_answer(_Completion.model_validate(completion.model_dump(mode="json"))) + + +def _chat_stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> _Answer: + stream: Final = _openai_client(gateway).chat.completions.create( + model=str(body["model"]), + messages=[{"role": "user", "content": str(body["question"])}], + stream=True, + stream_options={"include_usage": True}, + extra_body={key: value for key, value in body.items() if key not in ("model", "question")}, + ) + return _streamed_answer(tuple(_Chunk.model_validate(chunk.model_dump(mode="json")) for chunk in stream)) + + +def _unindexed(blocks: Sequence[dict[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + return tuple({key: value for key, value in block.items() if key != "index"} for block in blocks) + + +def _json_question(model: str, marker: str, cache: Mapping[str, JsonValue] = _NO_CACHE) -> dict[str, JsonValue]: + return { + "model": model, + "question": f"How many people live in Tokyo, in millions? {marker}", + "response_format": _RESPONSE_FORMAT, + "thinking": _THINKING, + "max_tokens": 2048, + "cache": dict(cache), + } + + +def test_fake_streamed_json_schema_chat_answer_matches_the_non_streamed_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Tokyo proper has about 14 million people, marker {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker) + + answer: Final = _chat(gateway, question) + streamed: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == ["converse", "converse"] + assert (answer.reasoning, answer.finish_reason, answer.tool_calls) == (reasoning, "stop", ()) + assert _JSON_OBJECT.validate_json(answer.content) == _ANSWER + assert answer.thinking_blocks == ( + {"type": "thinking", "thinking": reasoning, "signature": _SIGNATURE}, + {"type": "redacted_thinking", "data": _NATIVE_REDACTED}, + ) + assert (streamed.reasoning, streamed.finish_reason, streamed.tool_calls) == (reasoning, "stop", ()) + assert _JSON_OBJECT.validate_json(streamed.content) == _ANSWER + assert _unindexed(streamed.thinking_blocks) == answer.thinking_blocks + assert streamed.usage == answer.usage + assert answer.usage is not None and answer.usage.completion_tokens_details is not None + assert answer.usage.completion_tokens_details.reasoning_tokens == _THINKING_TOKENS + + replayed: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "thinking": _THINKING, + "max_tokens": 2048, + "cache": _NO_CACHE, + "messages": [ + {"role": "user", "content": str(question["question"])}, + { + "role": "assistant", + "content": streamed.content, + "thinking_blocks": list(streamed.thinking_blocks), + }, + {"role": "user", "content": "And Osaka?"}, + ], + }, + ) + assert replayed.status_code == 200, replayed.text + sent: Final = _ConverseRequest.model_validate_json(wire.drain()[-1].body) + assert next(message for message in sent.messages if message.role == "assistant").content == ( + {"reasoningContent": {"reasoningText": {"text": reasoning, "signature": _SIGNATURE}}}, + {"reasoningContent": {"redactedContent": _CONVERSE_REDACTED}}, + {"text": streamed.content}, + ) + + +def _spend_rows(model: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,) + ), + lambda found: len(found) >= count, + seconds=70, + ) + + +def test_fake_streamed_cached_prompt_is_billed_like_the_non_streamed_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_converse_peer(f"Billing check {marker}.")) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker) + _chat(gateway, question) + _chat_stream(gateway, question) + + rows: Final = _spend_rows(model, 2) + assert _converse_targets(wire) == ["converse", "converse"] + assert len(rows) == 2, rows + first, second = rows + assert (first["prompt_tokens"], first["completion_tokens"]) == ( + second["prompt_tokens"], + second["completion_tokens"], + ) + assert float(str(first["spend"])) > 0 + assert math.isclose(float(str(second["spend"])), float(str(first["spend"])), rel_tol=1e-9), rows + + +def _cache_holds(needle: str) -> bool: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as store: + return any(needle.encode() in _bytes(store.get(key)) for key in store.scan_iter(count=1000, _type="STRING")) + + +def _bytes(value: object) -> bytes: + return value if isinstance(value, bytes) else b"" + + +def test_cached_fake_stream_replays_the_answer_without_reaching_bedrock(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Cached reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + question: Final = _json_question(model, marker, cache={}) + + first: Final = _chat_stream(gateway, question) + assert first.reasoning == reasoning + eventually(lambda: _cache_holds(marker), bool, seconds=30) + assert _converse_targets(wire) == ["converse"] + + replay: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == [] + assert (replay.content, replay.finish_reason, replay.tool_calls) == (first.content, "stop", ()) + assert _unindexed(replay.thinking_blocks) in ((), _unindexed(first.thinking_blocks)) + + +def _responses_reasoning_and_answer( + events: Sequence[dict[str, JsonValue]], +) -> tuple[list[tuple[str, str, JsonValue]], str, str]: + lifecycle: Final = [ + (str(event["type"]), str(_JSON_OBJECT.validate_python(event["item"])["type"]), event["output_index"]) + for event in events + if event["type"] in ("response.output_item.added", "response.output_item.done") + ] + summary: Final = "".join( + str(event["delta"]) for event in events if event["type"] == "response.reasoning_summary_text.delta" + ) + text: Final = "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") + return lifecycle, summary, text + + +_REASONING_THEN_MESSAGE: Final = [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", "message", 1), + ("response.output_item.done", "message", 1), +] + + +def _raw_responses_events(gateway: Gateway, body: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + with gateway.client.stream( + "POST", + "/v1/responses", + json={**body, "stream": True, "store": False, "cache": _NO_CACHE}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + return [ + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: {") + ] + + +def _completed_usage(events: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + completed: Final = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1, [event["type"] for event in events] + return _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_python(completed[0]["response"])["usage"]) + + +def test_fake_streamed_json_schema_responses_stream_closes_the_reasoning_before_the_answer(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Responses reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + events: Final = _raw_responses_events( + gateway, + { + "model": model, + "input": f"How many people live in Tokyo, in millions? {marker}", + "reasoning": {"effort": "high"}, + "max_output_tokens": 2048, + "text": {"format": {"type": "json_schema", "name": "city_fact", "schema": _SCHEMA}}, + }, + ) + + lifecycle, summary, text = _responses_reasoning_and_answer(events) + assert _converse_targets(wire) == ["converse"] + assert lifecycle == _REASONING_THEN_MESSAGE, [event["type"] for event in events] + assert summary == reasoning + assert _JSON_OBJECT.validate_json(text) == _ANSWER + usage: Final = _completed_usage(events) + assert _JSON_OBJECT.validate_python(usage["input_tokens_details"])["cached_tokens"] == 50 + assert _JSON_OBJECT.validate_python(usage["output_tokens_details"])["reasoning_tokens"] == _THINKING_TOKENS + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _block_fields(block: anthropic.types.ContentBlock) -> tuple[str, str, str]: + match block.type: + case "thinking": + return (block.type, block.thinking, block.signature) + case "redacted_thinking": + return (block.type, block.data, "") + case "text": + return (block.type, block.text, "") + case _: + return (block.type, "", "") + + +def _message_blocks(message: anthropic.types.Message) -> list[tuple[str, str, str]]: + return [_block_fields(block) for block in message.content] + + +def test_fake_streamed_json_schema_messages_stream_opens_one_block_per_thinking_block(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Messages reasoning {marker}." + with wire_server(_converse_peer(reasoning)) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + with _anthropic_client(gateway).messages.stream( + model=model, + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[{"role": "user", "content": f"How many people live in Tokyo, in millions? {marker}"}], + extra_body={"output_format": {"type": "json_schema", "schema": _SCHEMA}, "cache": _NO_CACHE}, + ) as stream: + final: Final = stream.get_final_message() + + assert _converse_targets(wire) == ["converse"] + blocks: Final = _message_blocks(final) + assert blocks[:2] == [ + ("thinking", reasoning, _SIGNATURE), + ("redacted_thinking", _NATIVE_REDACTED, ""), + ], blocks + assert [block[0] for block in blocks[2:]] == ["text"], blocks + assert _JSON_OBJECT.validate_json(blocks[2][1]) == _ANSWER + assert final.stop_reason == "end_turn" + + +def _frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _converse_stream_with_thinking(reasoning_parts: Sequence[str], answer: str) -> Reply: + return Reply( + content_type=_EVENT_STREAM, + body=b"".join( + ( + _frame("messageStart", {"role": "assistant"}), + *( + _frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"text": part}}, "contentBlockIndex": 0}, + ) + for part in reasoning_parts + ), + _frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"signature": _SIGNATURE}}, "contentBlockIndex": 0}, + ), + _frame("contentBlockStop", {"contentBlockIndex": 0}), + _frame("contentBlockDelta", {"delta": {"text": answer}, "contentBlockIndex": 1}), + _frame("contentBlockStop", {"contentBlockIndex": 1}), + _frame("messageStop", {"stopReason": "end_turn"}), + _frame("metadata", {"usage": {"inputTokens": 11, "outputTokens": 9, "totalTokens": 20}}), + ) + ), + ) + + +def test_native_converse_stream_thinking_reaches_the_anthropic_sdk_once_with_its_signature(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + parts: Final = ("Native thinking ", f"marker {marker}.") + answer: Final = "It is sunny in Paris." + + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse-stream"), request.target + return _converse_stream_with_thinking(parts, answer) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _claude_model(scenario, wire) + with _anthropic_client(gateway).messages.stream( + model=model, + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[{"role": "user", "content": f"Weather in Paris? {marker}"}], + extra_body={"cache": _NO_CACHE}, + ) as stream: + final: Final = stream.get_final_message() + + assert _converse_targets(wire) == ["converse-stream"] + assert _message_blocks(final) == [("thinking", "".join(parts), _SIGNATURE), ("text", answer, "")] + + +def _llama_peer(content: list[JsonValue], stop_reason: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert unquote(request.target).endswith("/converse"), request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": content}}, + "stopReason": stop_reason, + "usage": {"inputTokens": 175, "outputTokens": 23, "totalTokens": 198}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + return respond + + +_PARIS: Final[dict[str, JsonValue]] = {"city": "Paris"} + + +@pytest.mark.parametrize( + "content, stop_reason, expected_calls", + [ + pytest.param( + [ + {"toolUse": {"toolUseId": "tooluse_weather", "name": "get_weather", "input": _PARIS}}, + {"toolUse": {"toolUseId": "tooluse_time", "name": "get_time", "input": _PARIS}}, + ], + "tool_use", + (("get_weather", _PARIS), ("get_time", _PARIS)), + id="two-tool-use-blocks", + ), + pytest.param( + [{"text": json.dumps({"type": "function", "name": "get_weather", "parameters": _PARIS})}], + "end_turn", + (("get_weather", _PARIS),), + id="tool-call-written-as-text", + ), + ], +) +def test_fake_streamed_llama_tool_calls_match_the_non_streamed_answer( + gateway: Gateway, + content: list[JsonValue], + stop_reason: str, + expected_calls: tuple[tuple[str, JsonValue], ...], +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_llama_peer(content, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/converse/{_LLAMA_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + question: Final[dict[str, JsonValue]] = { + "model": model, + "question": f"Weather and time in Paris? {marker}", + "tools": [_tool("get_weather"), _tool("get_time")], + "cache": _NO_CACHE, + } + + answer: Final = _chat(gateway, question) + streamed: Final = _chat_stream(gateway, question) + + assert _converse_targets(wire) == ["converse", "converse"] + assert (answer.tool_calls, answer.finish_reason) == (expected_calls, "tool_calls") + assert (streamed.tool_calls, streamed.finish_reason, streamed.content) == ( + answer.tool_calls, + "tool_calls", + answer.content, + ) + + +def _azure_o_series_reply() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "o4-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_PARIS)}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": json.dumps(_PARIS)}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, + } + ).encode() + + +def test_azure_o_series_fake_stream_keeps_every_tool_call_and_token_detail(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(lambda _: Reply(body=_azure_o_series_reply())) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/o_series/{_AZURE_DEPLOYMENT}", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version=_AZURE_API_VERSION, + ) + streamed: Final = _chat_stream( + gateway, + { + "model": model, + "question": f"Weather and time in Paris? {marker}", + "tools": [_tool("get_weather"), _tool("get_time")], + "cache": _NO_CACHE, + }, + ) + + received: Final = wire.drain() + assert [unquote(request.target).split("?")[0] for request in received] == [ + f"/openai/deployments/{_AZURE_DEPLOYMENT}/chat/completions" + ] + assert _JSON_OBJECT.validate_json(received[0].body).get("stream") is not True + assert (streamed.tool_calls, streamed.finish_reason) == ( + (("get_weather", _PARIS), ("get_time", _PARIS)), + "tool_calls", + ) + assert streamed.usage is not None + assert streamed.usage.prompt_tokens_details == _PromptDetails(cached_tokens=64) + assert streamed.usage.completion_tokens_details == _CompletionDetails(reasoning_tokens=1920) + + +def _vllm_frame(identity: str, delta: Mapping[str, JsonValue], finish_reason: str | None = None) -> bytes: + usage: Final = ( + {"usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}} if finish_reason else {} + ) + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _VLLM_BACKEND, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + **usage, + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _vllm_sse(deltas: Sequence[Mapping[str, JsonValue]], finish_reason: str) -> Reply: + identity: Final = f"chatcmpl-audit-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + _vllm_frame(identity, {"role": "assistant", "content": ""}), + *(_vllm_frame(identity, delta) for delta in deltas), + _vllm_frame(identity, {}, finish_reason), + b"data: [DONE]\n\n", + ), + ) + + +def _is_discovery_probe(request: Request) -> bool: + return (request.method, request.target) == _DISCOVERY_PROBE + + +@contextmanager +def _vllm_server(reply: Reply) -> Generator[Wire, None, None]: + with wire_server( + lambda request: Reply(body=b'{"object":"list","data":[]}') if _is_discovery_probe(request) else reply + ) as wire: + yield wire + + +def _streamed_chats(wire: Wire) -> list[JsonValue]: + posted: Final = [request for request in wire.drain() if not _is_discovery_probe(request)] + assert [(request.method, request.target) for request in posted] == [("POST", "/v1/chat/completions")] + return [_JSON_OBJECT.validate_json(request.body).get("stream") for request in posted] + + +@pytest.mark.parametrize( + "deltas, merged", + [ + pytest.param( + ({"reasoning_content": "R1 "}, {"reasoning_content": "R2 ", "content": "A1 "}, {"content": "A2"}), + "R1 R2 A1 A2", + id="answer-starts-in-the-last-reasoning-chunk", + ), + pytest.param( + ({"reasoning_content": "R ", "content": "A "}, {"content": "B"}), + "R A B", + id="reasoning-and-answer-in-the-first-chunk", + ), + pytest.param( + ({"reasoning_content": "R1 "}, {"reasoning_content": "R2 "}, {"content": "A1 "}, {"content": "A2"}), + "R1 R2 A1 A2", + id="reasoning-and-answer-in-separate-chunks", + ), + ], +) +def test_merged_reasoning_content_closes_the_think_tag_before_the_answer( + gateway: Gateway, deltas: tuple[dict[str, JsonValue], ...], merged: str +) -> None: + with _vllm_server(_vllm_sse(deltas, "stop")) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"hosted_vllm/{_VLLM_BACKEND}", + api_base=wire.url + "/v1", + api_key=_VLLM_KEY, + merge_reasoning_content_in_choices=True, + ) + streamed: Final = _chat_stream( + gateway, {"model": model, "question": f"Hello {uuid.uuid4().hex}", "cache": _NO_CACHE} + ) + + assert _streamed_chats(wire) == [True] + assert (streamed.content, streamed.reasoning, streamed.finish_reason) == (merged, "", "stop") + + +_VLLM_TOOL_CALL: Final[dict[str, JsonValue]] = { + "tool_calls": [ + { + "index": 0, + "id": "call_bridge_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps(_PARIS)}, + } + ] +} + + +@pytest.mark.parametrize( + "answer_delta, finish_reason, answer_item", + [ + pytest.param({"content": "It is sunny."}, "stop", "message", id="text"), + pytest.param(_VLLM_TOOL_CALL, "tool_calls", "function_call", id="tool-call"), + ], +) +def test_responses_stream_closes_the_reasoning_item_before_an_answer_from_the_same_chunk( + gateway: Gateway, answer_delta: dict[str, JsonValue], finish_reason: str, answer_item: str +) -> None: + marker: Final = uuid.uuid4().hex + reasoning: Final = f"Bridge reasoning {marker}." + reply: Final = _vllm_sse(({"reasoning_content": reasoning, **answer_delta},), finish_reason) + with _vllm_server(reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"hosted_vllm/{_VLLM_BACKEND}", + api_base=wire.url + "/v1", + api_key=_VLLM_KEY, + use_chat_completions_api=True, + ) + events: Final = _raw_responses_events( + gateway, + { + "model": model, + "input": f"Weather in Paris? {marker}", + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ], + }, + ) + + lifecycle, summary, text = _responses_reasoning_and_answer(events) + assert _streamed_chats(wire) == [True] + assert lifecycle == [ + ("response.output_item.added", "reasoning", 0), + ("response.output_item.done", "reasoning", 0), + ("response.output_item.added", answer_item, 1), + ("response.output_item.done", answer_item, 1), + ], [event["type"] for event in events] + assert summary == reasoning + assert text == str(answer_delta.get("content", "")) 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/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index 3fba69b39b7..9e8f456df49 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -109,7 +109,9 @@ class FallbackProxy: spare_port: int -def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Path: +def _fallback_config( + directory: Path, primary_port: int, spare_port: int, extra_router_settings: Mapping[str, JsonValue] +) -> Path: base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) deployments: Final = [ { @@ -122,23 +124,37 @@ def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Pat } for name, port in ((_PRIMARY, primary_port), (_SPARE, spare_port)) ] - router_settings: Final = {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{_PRIMARY: [_SPARE]}]} + router_settings: Final = { + "num_retries": 0, + "disable_cooldowns": True, + "fallbacks": [{_PRIMARY: [_SPARE]}], + **extra_router_settings, + } path: Final = directory / "bridged-fallbacks.yaml" path.write_text(yaml.safe_dump({**base, "model_list": deployments, "router_settings": router_settings})) return path -@pytest.fixture(scope="module") -def fallback_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: +def _fallback_proxy( + tmp_path_factory: pytest.TempPathFactory, extra_router_settings: Mapping[str, JsonValue] +) -> Iterator[FallbackProxy]: directory: Final = tmp_path_factory.mktemp("bridged-fallbacks") primary_port, spare_port = allocate_port(), allocate_port() - with ( - gateway_from_environment() as shared, - owned_proxy(shared, directory, {}, config=_fallback_config(directory, primary_port, spare_port)) as owned, - ): + config: Final = _fallback_config(directory, primary_port, spare_port, extra_router_settings) + with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=config) as owned: yield FallbackProxy(owned, primary_port, spare_port) +@pytest.fixture(scope="module") +def fallback_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: + yield from _fallback_proxy(tmp_path_factory, {}) + + +@pytest.fixture(scope="module") +def bad_request_opt_out_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: + yield from _fallback_proxy(tmp_path_factory, {"fallback_on_unsupported_params": False}) + + def test_bridged_in_stream_rate_limit_falls_back_to_the_healthy_deployment(fallback_proxy: FallbackProxy) -> None: identity: Final = "resp_" + uuid.uuid4().hex with ( @@ -189,9 +205,7 @@ def test_bridged_top_level_rate_limit_falls_back_to_the_healthy_deployment(fallb _INVALID_PROMPT: Final = "Invalid prompt: your prompt was flagged." - - -@pytest.mark.parametrize( +_INVALID_PROMPT_REJECTIONS: Final = pytest.mark.parametrize( "rejection", [ pytest.param( @@ -201,7 +215,10 @@ _INVALID_PROMPT: Final = "Invalid prompt: your prompt was flagged." pytest.param(openai_error_event("invalid_prompt", _INVALID_PROMPT), id="top_level"), ], ) -def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallback( + + +@_INVALID_PROMPT_REJECTIONS +def test_bridged_invalid_prompt_before_any_token_falls_back_to_the_healthy_deployment( fallback_proxy: FallbackProxy, rejection: Mapping[str, JsonValue] ) -> None: identity: Final = "resp_" + uuid.uuid4().hex @@ -214,6 +231,29 @@ def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallbac ) as spare, ): response: Final = _chat(fallback_proxy.gateway, chat_body(_PRIMARY, identity)) + assert response.status_code == 200, response.text + assert chat_content(data_frames(response.text)) == "fallback answer", response.text + assert len(primary.drain()) == 1 and len(spare.drain()) == 1 + row: Final = spend_row(response.headers["x-litellm-call-id"]) + assert row["status"] == "success" and row["model_group"] == _SPARE, row + + +@_INVALID_PROMPT_REJECTIONS +def test_bridged_invalid_prompt_answers_a_bad_request_without_trying_the_fallback_when_bad_requests_opt_out( + bad_request_opt_out_proxy: FallbackProxy, rejection: Mapping[str, JsonValue] +) -> None: + identity: Final = "resp_" + uuid.uuid4().hex + with ( + wire_server( + serve((frame(created(identity)), frame(rejection)), AZURE_TARGET), + port=bad_request_opt_out_proxy.primary_port, + ) as primary, + wire_server( + serve(healthy_stream(identity, "fallback answer"), AZURE_TARGET), + port=bad_request_opt_out_proxy.spare_port, + ) as spare, + ): + response: Final = _chat(bad_request_opt_out_proxy.gateway, chat_body(_PRIMARY, identity)) assert response.status_code == 400, response.text message: Final = string_value(error_body(response)["message"]) assert message.startswith("litellm.BadRequestError: ") and _INVALID_PROMPT in message, message diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/test_litellm_rust/chat_completions/test_dispatch.py similarity index 100% rename from tests/unit/chat_completions/test_dispatch.py rename to tests/test_litellm_rust/chat_completions/test_dispatch.py diff --git a/tests/unit/chat_completions/__init__.py b/tests/test_litellm_rust/embeddings/__init__.py similarity index 100% rename from tests/unit/chat_completions/__init__.py rename to tests/test_litellm_rust/embeddings/__init__.py diff --git a/tests/unit/embeddings/test_dispatch.py b/tests/test_litellm_rust/embeddings/test_dispatch.py similarity index 100% rename from tests/unit/embeddings/test_dispatch.py rename to tests/test_litellm_rust/embeddings/test_dispatch.py diff --git a/tests/unit/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py similarity index 100% rename from tests/unit/messages/test_dispatch.py rename to tests/test_litellm_rust/messages/test_dispatch.py diff --git a/tests/unit/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_public_dispatch.py similarity index 100% rename from tests/unit/ocr/test_dispatch.py rename to tests/test_litellm_rust/ocr/test_public_dispatch.py diff --git a/tests/unit/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py similarity index 100% rename from tests/unit/responses/test_dispatch.py rename to tests/test_litellm_rust/responses/test_dispatch.py diff --git a/tests/unit/rust_bridge/test_dispatch.py b/tests/test_litellm_rust/test_dispatch.py similarity index 100% rename from tests/unit/rust_bridge/test_dispatch.py rename to tests/test_litellm_rust/test_dispatch.py 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/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index a9541f9fcb8..58549a028ea 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -52,7 +52,8 @@ import json import httpx import respx from fastapi.testclient import TestClient -from litellm._internal_context import current_service_target, in_post_response_phase +from contextlib import nullcontext +from litellm._internal_context import current_service_target, in_post_response_phase, internal_sub_call from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -3046,3 +3047,57 @@ async def test_async_get_cache_answers_400_for_a_single_object_embedding_input() call_type=CallTypes.aembedding.value, kwargs={"model": model, "custom_llm_provider": "gemini", "input": clip_block, "caching": True}, ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("internal", [False, True]) +async def test_a_cache_hit_inside_an_internal_sub_call_is_free_and_fires_no_callbacks_of_its_own( + monkeypatch: pytest.MonkeyPatch, internal: bool +): + cache: Final = Cache(type="local") + monkeypatch.setattr(litellm, "cache", cache) + kwargs: Final = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True} + cache.add_cache( + result=litellm.ModelResponse( + model="gpt-5.4", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))], + usage=litellm.Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), + ), + **kwargs, + ) + async_logging: Final = _build_logging_obj(CallTypes.acompletion.value, stream=False) + async_logging.async_success_handler = AsyncMock() + async_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() + sync_logging: Final = _build_logging_obj(CallTypes.completion.value, stream=False) + sync_logging.handle_sync_success_callbacks_for_async_calls = MagicMock() + + with internal_sub_call() if internal else nullcontext(): + async_hit: Final = await LLMCachingHandler( + original_function=litellm.acompletion, request_kwargs=kwargs, start_time=_FIXED_START + ).async_get_cache( + model="gpt-5.4", + original_function=litellm.acompletion, + logging_obj=async_logging, + start_time=_FIXED_START, + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + args=(), + ) + sync_hit: Final = LLMCachingHandler( + original_function=completion, request_kwargs=kwargs, start_time=_FIXED_START + ).sync_get_cache( + model="gpt-5.4", + original_function=completion, + logging_obj=sync_logging, + start_time=_FIXED_START, + call_type=CallTypes.completion.value, + kwargs=kwargs, + ) + + assert async_hit is not None and async_hit.cached_result is not None + assert sync_hit.cached_result is not None + assert async_hit.cached_result._hidden_params.get("response_cost") == 0.0 + assert sync_hit.cached_result._hidden_params.get("response_cost") == 0.0 + assert async_logging.async_success_handler.called is not internal + assert async_logging.handle_sync_success_callbacks_for_async_calls.called is not internal + assert sync_logging.handle_sync_success_callbacks_for_async_calls.called is not internal diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 8b1a536d8e3..46b768e4924 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -120,6 +120,7 @@ _PROVIDERS: Final[tuple[tuple[str, str, str, str], ...]] = ( "pplx-decider-v1-27b", ), ("typesafe", "typesafe/jev-1.13", "https://api.typesafe.ai/v1/systemone", "jev-1.13"), + ("nace", "nace/drex-v2.0", "https://console.nace.ai/v1/systemone", "drex-v2.0"), ( "openrouter", "openrouter/typesafe/jev-1.13", @@ -3345,3 +3346,137 @@ def test_dropping_every_question_answers_400_before_http(respx_mock: respx.MockR assert caught.value.status_code == 400 assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("openai_format", (False, True)) +@pytest.mark.parametrize( + ("provider", "model", "url", "key_env"), + ( + ("nace", "drex-test-alias", "https://console.nace.ai/v1/systemone", "NACE_API_KEY"), + ("openrouter", "nace/drex-test-alias", "https://openrouter.ai/api/alpha/decisions", "OPENROUTER_API_KEY"), + ( + "fireworks_ai", + "accounts/test/models/drex-test-alias", + "https://api.fireworks.ai/inference/v1/systemone", + "FIREWORKS_API_KEY", + ), + ), +) +async def test_decision_hosts_preserve_drex_model_routing_credentials_and_returned_model_pricing( + provider: str, + model: str, + url: str, + key_env: str, + openai_format: bool, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv(key_env, "host-test-key") + monkeypatch.delenv("NACE_API_BASE", raising=False) + monkeypatch.delenv("NACE_BASE_URL", raising=False) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + monkeypatch.delenv("FIREWORKS_API_BASE", raising=False) + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/{model}", + { + "litellm_provider": provider, + "mode": "evaluation", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0, + }, + ) + monkeypatch.setitem( + litellm.model_cost, + f"{provider}/drex-test-resolved", + { + "litellm_provider": provider, + "mode": "evaluation", + "input_cost_per_token": 0.002, + "output_cost_per_token": 0, + }, + ) + payload: Final = {**_RESPONSE, "model": "drex-test-resolved", "request_id": "request-test"} + upstream: Final = respx_mock.post(url).respond(json=payload) + if openai_format: + response: Final = await litellm.adecisions( + model=f"{provider}/{model}", + input="damaged order", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + ) + assert isinstance(response, OpenAIDecisionResponse) + assert response.answers[0].model_dump() == {"type": "predicate", "name": "is_defect", "probability": 0.9} + else: + systemone_response: Final = await litellm.asystemone( + model=f"{provider}/{model}", + state="damaged order", + questions={"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + ) + assert systemone_response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert litellm.completion_cost(completion_response=systemone_response, call_type="asystemone") == pytest.approx( + _INPUT_TOKENS * 0.002 + ) + assert upstream.call_count == 1 + assert upstream.calls[0].request.headers["authorization"] == "Bearer host-test-key" + assert json.loads(upstream.calls[0].request.content) == { + "model": model, + "state": "damaged order", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + + +@pytest.mark.parametrize("base_env", ("NACE_API_BASE", "NACE_BASE_URL")) +def test_nace_sync_call_uses_environment_base_and_explicit_credentials( + base_env: str, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.setenv(base_env, "https://nace.example/v1/") + monkeypatch.setenv("NACE_API_KEY", "environment-key") + upstream: Final = respx_mock.post("https://nace.example/v1/systemone").respond(json=_RESPONSE) + response: Final = litellm.systemone( + model="nace/drex-v2.0", + state={"message": "damaged order"}, + api_key="explicit-key", + questions={"is_defect": {"type": "noul", "criteria": {"true": "damaged", "false": "intact"}}}, + ) + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert upstream.calls[0].request.headers["authorization"] == "Bearer explicit-key" + assert json.loads(upstream.calls[0].request.content) == { + "model": "drex-v2.0", + "state": {"message": "damaged order"}, + "questions": {"is_defect": {"type": "noul", "criteria": {"true": "damaged", "false": "intact"}}}, + } + + +@pytest.mark.parametrize( + ("status", "exception"), + ( + (401, litellm.AuthenticationError), + (402, litellm.PaymentRequiredError), + (422, litellm.BadRequestError), + (429, litellm.RateLimitError), + (529, litellm.APIError), + ), +) +def test_nace_upstream_failures_preserve_status_and_message( + status: int, + exception: type[Exception], + respx_mock: respx.MockRouter, +) -> None: + upstream: Final = respx_mock.post("https://console.nace.ai/v1/systemone").respond( + status, + json={"error": {"message": "decision rejected"}, "request_id": "request-test"}, + headers={"retry-after": "1"}, + ) + with pytest.raises(exception, match="decision rejected") as error: + litellm.systemone( + model="nace/drex-v2.0", + state="text", + api_key="test-key", + questions={"a": {"type": "noul", "instructions": "Is it urgent?"}}, + num_retries=0, + ) + assert getattr(error.value, "status_code") == status + assert upstream.call_count == 1 diff --git a/tests/unit/embeddings/__init__.py b/tests/unit/embeddings/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 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/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py index 0cbcf1b7ee1..9e114e3b0b7 100644 --- a/tests/unit/integration_support/test_upstream.py +++ b/tests/unit/integration_support/test_upstream.py @@ -21,7 +21,11 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},)) provider.scenario_store.put("owned-scenario", response) incoming: Final[deque[Message]] = deque( - ({"type": "websocket.connect"}, {"type": "websocket.receive", "text": '{"type":"response.create"}'}) + ( + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": '{"type":"response.create"}'}, + {"type": "websocket.disconnect", "code": 1006}, + ) ) outgoing: Final[deque[Message]] = deque() 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/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py index 45b378d10fe..28d744cdb06 100644 --- a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py +++ b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py @@ -20,10 +20,14 @@ Tests cover: - llm_router unavailable / model_group missing / router raises → silent no-op. """ +from collections.abc import Iterator +from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from prometheus_client.samples import Sample from litellm.integrations.prometheus import PrometheusLogger @@ -293,3 +297,42 @@ class TestRouterFallbackDefensivePaths: ) prometheus_logger.litellm_remaining_tokens_metric.labels.assert_not_called() + + +def _registry_samples() -> Iterator[Sample]: + for metric in REGISTRY.collect(): + yield from metric.samples + + +def _sample_values(sample_name: str) -> tuple[float, ...]: + return tuple(sample.value for sample in _registry_samples() if sample.name == sample_name) + + +def test_unparseable_provider_ratelimit_header_skips_only_its_own_gauge( + prometheus_logger: PrometheusLogger, +): + payload: Final = _build_payload( + model_group="groq-group", + custom_llm_provider="groq", + additional_headers={ + "x_ratelimit_remaining_requests": 0, + "x_ratelimit_remaining_tokens": "856, abc", + }, + ) + request_kwargs: Final = { + "litellm_params": {"custom_llm_provider": "groq", "metadata": {}}, + "model": "llama-3.1-8b-instant", + "standard_logging_object": payload, + } + + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs=request_kwargs, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1), + enum_values=_enum_values(model_group="groq-group"), + output_tokens=2.0, + ) + + assert _sample_values("litellm_remaining_requests_metric") == (0.0,) + assert _sample_values("litellm_remaining_tokens_metric") == () + assert _sample_values("litellm_deployment_success_responses_total") == (1.0,) 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_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index e891706df13..3cbc3591009 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -25,13 +25,13 @@ from typing import Any, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from openai._legacy_response import LegacyAPIResponse import litellm -from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, ) @@ -45,9 +45,9 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, ) from litellm.types.utils import ( + ChatCompletionMessageToolCall, Choices, Function, - ChatCompletionMessageToolCall, Message, ModelResponse, ) @@ -142,7 +142,7 @@ def _raw_response_for(model_response: ModelResponse) -> MagicMock: (an object exposing `.headers` and `.parse()` -> something with model_dump).""" parsed = MagicMock() parsed.model_dump.return_value = model_response.model_dump() - raw = MagicMock() + raw = MagicMock(spec=LegacyAPIResponse) raw.headers = {} raw.parse.return_value = parsed return raw @@ -171,7 +171,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal ] ) mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create = create + mock_client.post = create response = await litellm.acompletion( model="openai/gpt-4o-mini", diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 9a7f96000e6..c3a2f032651 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -1,6 +1,8 @@ """Tests for litellm_core_utils.core_helpers module.""" import logging +import threading +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, + get_litellm_metadata_from_kwargs, get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, get_provider_response_headers_from_hidden_params, @@ -587,3 +590,55 @@ class TestProviderResponseHeadersInHiddenParams: assert get_provider_response_headers_from_hidden_params(sibling) is None assert "additional_headers" not in sibling._hidden_params + + +def test_spend_metadata_merge_survives_the_logging_thread_inserting_into_the_request_metadata() -> None: + metadata: Final[dict[str, object]] = {} + + def logging_thread_records_hidden_params() -> None: + metadata["hidden_params"] = {"response_cost": 0.1} + + class KeyCheckedWhileTheLoggingThreadRuns(str): + def __contains__(self, key: str) -> bool: + writer: Final = threading.Thread(target=logging_thread_records_hidden_params) + writer.start() + writer.join() + return super().__contains__(key) + + metadata[KeyCheckedWhileTheLoggingThreadRuns("user_api_key_alias")] = "alias" + metadata["user_api_key_hash"] = "hash" + metadata["requester_ip_address"] = "10.0.0.1" + litellm_metadata: Final[dict[str, object]] = {"user_api_key_alias": "stale", "model_group": "claude"} + + merged: Final = get_litellm_metadata_from_kwargs( + {"litellm_params": {"metadata": metadata, "litellm_metadata": litellm_metadata}} + ) + + assert "hidden_params" in metadata + assert merged == {"user_api_key_alias": "alias", "user_api_key_hash": "hash", "model_group": "claude"} + + +def test_spend_metadata_merge_reads_request_metadata_that_is_a_read_only_mapping() -> None: + class ReadOnlyMetadata(Mapping[str, object]): + def __init__(self, values: dict[str, object]) -> None: + self._values: Final = values + + def __getitem__(self, key: str) -> object: + return self._values[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + merged: Final = get_litellm_metadata_from_kwargs( + { + "litellm_params": { + "metadata": ReadOnlyMetadata({"user_api_key_hash": "hash", "requester_ip_address": "10.0.0.1"}), + "litellm_metadata": {"model_group": "claude"}, + } + } + ) + + assert merged == {"user_api_key_hash": "hash", "model_group": "claude"} diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 02ebae6dcf6..b24348cfaa8 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -2462,7 +2462,6 @@ def test_openai_gateway_timeout_error(): Test that the OpenAI gateway timeout error is raised """ openai_client = OpenAI() - mapped_target = openai_client.chat.completions.with_raw_response # type: ignore def _return_exception(*args, **kwargs): @@ -2498,8 +2497,8 @@ def test_openai_gateway_timeout_error(): with pytest.raises(litellm.Timeout) as exc_info: with patch.object( - mapped_target, - "create", + openai_client, + "post", side_effect=_return_exception, ): litellm.completion( 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_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index 9d1e1e8dceb..53b6c0370fa 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params, get_hidden_params_storage, get_or_create_hidden_params, + served_from_cache, set_hidden_param, set_hidden_params, ) @@ -250,3 +251,24 @@ def test_set_hidden_params_replaces_frozen_decisions_response_private_attr() -> assert response.hidden_params is replacement assert response._hidden_params is replacement + + +@pytest.mark.parametrize( + "hidden_params, expected", + [ + ({"cache_hit": True}, True), + ({"cache_hit": None}, False), + ({"cache_hit": False}, False), + ({"cache_hit": "True"}, False), + ({}, False), + ], +) +def test_served_from_cache_is_true_only_for_a_response_the_cache_stamped( + hidden_params: dict[str, object], expected: bool +) -> None: + model_response: Final = ModelResponse() + set_hidden_params(model_response, dict(hidden_params)) + + assert served_from_cache(model_response) is expected + assert served_from_cache({HIDDEN_PARAMS_ATTR: dict(hidden_params)}) is expected + assert served_from_cache(object()) is False diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index df6593d29b1..31eaf12a186 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -823,6 +823,20 @@ class TestZeroCostDiagnostic: assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None assert self._zero_cost_warnings(caplog) == [] + def test_response_served_from_the_cache_at_no_charge_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=0.0, cache_hit=True, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj.response_cost_calculator(result=response) + + assert cost == 0.0 + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + def test_unbilled_read_route_with_usage_stays_silent( self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture ) -> None: @@ -5363,6 +5377,63 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +@pytest.mark.parametrize( + "hidden_params, passed_cache_hit, expected_cache_hit, expected_cost, expected_unique_id", + [ + ({"cache_hit": True, "response_cost": 0.0}, None, True, 0.0, True), + ({"cache_hit": True, "response_cost": 0.002}, None, None, 0.002, True), + ({"cache_hit": True, "response_cost": 0.0}, False, False, 0.0, True), + ({"response_cost": 0.0}, None, None, 0.0, False), + ({}, None, None, None, False), + ], +) +def test_success_handler_logs_a_response_served_from_the_cache_under_a_unique_id( + hidden_params: dict[str, object], + passed_cache_hit: bool | None, + expected_cache_hit: bool | None, + expected_cost: float | None, + expected_unique_id: bool, +): + """A parent call whose response came back from the response cache (an internal sub-call's hit) reuses the + id of the row that cached it, so its row needs a unique id, and it is a cache hit only when it cost nothing.""" + from litellm.litellm_core_utils.hidden_params import set_hidden_params + + logging_obj: Final = LitellmLogging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="parent-call-id", + function_id="parent-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="gpt-5.6", custom_llm_provider="openai" + ) + response: Final = ModelResponse( + id="chatcmpl-cached-src", + model="gpt-5.6", + usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), + ) + set_hidden_params(response, dict(hidden_params)) + + logging_obj._success_handler_helper_fn( + result=response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + cache_hit=passed_cache_hit, + ) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert logging_obj.model_call_details["cache_hit"] is expected_cache_hit + assert payload["cache_hit"] is expected_cache_hit + if expected_cost is not None: + assert payload["response_cost"] == expected_cost + assert payload["id"].startswith("chatcmpl-cached-src_cache_hit") is expected_unique_id + assert (payload["id"] == "chatcmpl-cached-src") is not expected_unique_id + assert payload["saved_cache_cost"] == 0.0 or expected_cache_hit is True + + @pytest.mark.parametrize( "client_sent_oauth_token, custom_llm_provider, expected", [(True, "anthropic", True), (True, "bedrock", False), (False, "anthropic", False), (None, "anthropic", None)], @@ -6066,6 +6137,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_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 8de5e27d99c..7d85da3f756 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -13,6 +13,7 @@ from websockets.frames import Close import litellm from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -543,6 +544,32 @@ async def test_backend_to_client_stores_normalized_events_for_logging(): assert streaming.messages[0]["response"]["usage"]["total_tokens"] == 0 +@pytest.mark.asyncio +async def test_client_ack_messages_replays_the_opening_message_a_refused_bedrock_stream_left(): + opening_message: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is 17 * 3?"}]}, + } + ) + next_message: Final = json.dumps({"type": "response.create"}) + client_ws = MagicMock() + client_ws.scope = {"headers": [], BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: opening_message} + client_ws.receive_text = AsyncMock(side_effect=[next_message, Exception("connection closed")]) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in backend_ws.send.call_args_list] + assert sent_types == ["conversation.item.create", "response.create"] + assert json.loads(backend_ws.send.call_args_list[0].args[0])["item"]["content"][0]["text"] == "What is 17 * 3?" + assert BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY not in client_ws.scope + + @pytest.mark.asyncio async def test_client_ack_messages_keeps_beta_session_shape_for_beta_clients(): client_ws = MagicMock() @@ -2351,8 +2378,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value( injected_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(injected_turn_detection, dict) - assert injected_turn_detection["create_response"] is False + assert injected_turn_detection == {"type": "server_vad", "create_response": False} assert streaming._guardrail_turn_detection_update_sent is True @@ -2417,8 +2443,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act forwarded_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(forwarded_turn_detection, dict) - assert forwarded_turn_detection["create_response"] is False + assert forwarded_turn_detection == {"type": "server_vad", "create_response": False} @pytest.mark.asyncio @@ -3927,3 +3952,317 @@ async def test_transcription_session_guardrail_block_only_reports_violation( sent_to_backend: Final = [call.args[0] for call in backend_ws.send.await_args_list] assert sent_to_backend == [], sent_to_backend assert backend_ws.close.await_count == (1 if expect_session_closed else 0) + + +_GA_GUARDRAIL_GATE: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, +} + + +async def _forwarded_ga_session_updates( + monkeypatch: pytest.MonkeyPatch, + *client_sessions: dict[str, object], + event_normalizer: XAIRealtimeNormalizer | None = None, +) -> tuple[dict[str, object], ...]: + """Run a GA client's session.updates through a transcript-guardrailed proxy and return the forwarded sessions.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + *(json.dumps({"type": "session.update", "session": session}) for session in client_sessions), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + await RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer).client_ack_messages() + return tuple(json.loads(call.args[0])["session"] for call in backend_ws.send.await_args_list) + + +@pytest.mark.asyncio +async def test_first_ga_session_update_keeps_client_turn_detection_when_guardrail_disables_auto_response( + monkeypatch: pytest.MonkeyPatch, +): + """A transcript guardrail only flips create_response: the client's nested GA + turn_detection, its type included, must reach the backend intact, or OpenAI + rejects the whole update and the client's session config is lost.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + { + "type": "realtime", + "instructions": "hi", + "audio": { + "input": { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": True}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + }, + }, + ) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["instructions"] == "hi" + assert client_update["audio"]["input"] == { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": False}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + + +@pytest.mark.asyncio +async def test_first_ga_session_update_without_turn_detection_gets_a_typed_gate(monkeypatch: pytest.MonkeyPatch): + forwarded: Final = await _forwarded_ga_session_updates(monkeypatch, {"type": "realtime", "instructions": "hi"}) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + + +@pytest.mark.asyncio +async def test_later_ga_session_update_without_turn_detection_is_forwarded_without_one( + monkeypatch: pytest.MonkeyPatch, +): + """OpenAI keeps the earlier create_response=false for a field left unset, while a + bare create_response carries no type and gets the whole update rejected.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + {"type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": True}}}}, + {"type": "realtime", "instructions": "later"}, + ) + + gate, first, later = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert first["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + assert later == {"type": "realtime", "instructions": "later"} + + +_XAI_TYPED_VAD_SESSION: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 700}}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_sessions", + [ + pytest.param(({"type": "realtime", "instructions": "hi"},), id="no-turn-detection"), + pytest.param(({"type": "realtime", "audio": {"input": {"turn_detection": None}}},), id="push-to-talk"), + pytest.param((_XAI_TYPED_VAD_SESSION,), id="typed-server-vad"), + pytest.param( + ({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad" + ), + ], +) +async def test_xai_session_update_gate_leaves_out_the_vad_type( + monkeypatch: pytest.MonkeyPatch, client_sessions: tuple[dict[str, object], ...] +): + """xAI answers on its own under a typed turn_detection even with create_response off, + so every turn_detection the proxy forwards to xAI under the guardrail must carry no type.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, *client_sessions, event_normalizer=XAIRealtimeNormalizer() + ) + + assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [ + {"create_response": False} + ] * (len(client_sessions) + 1) + + +@pytest.mark.asyncio +async def test_xai_session_created_gate_leaves_out_the_vad_type(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=XAIRealtimeNormalizer()) + + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert sent == [ + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": {"create_response": False}}}}, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_turn_detection", + [ + pytest.param({"type": "semantic_vad", "eagerness": "low"}, id="semantic-vad"), + pytest.param({"type": "server_vad", "silence_duration_ms": 700}, id="tuned-server-vad"), + ], +) +async def test_guardrail_gate_goes_ahead_of_a_session_update_sent_before_session_created( + monkeypatch: pytest.MonkeyPatch, client_turn_detection: dict[str, object] +): + """OpenAI applies events in order, so the gate must precede a client update sent on connect: + the client's turn_detection then wins, and when OpenAI refuses that update the gate still holds.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": client_turn_detection}}}, + } + ), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[json.dumps({"type": "session.created", "session": {}}).encode(), ConnectionClosed(None, None)] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["session"]["audio"]["input"]["turn_detection"] for event in sent] == [ + {"type": "server_vad", "create_response": False}, + {**client_turn_detection, "create_response": False}, + ] + + +@pytest.mark.asyncio +async def test_client_never_sees_the_error_caused_by_the_guardrail_cancel(monkeypatch: pytest.MonkeyPatch): + """OpenAI answers the guardrail's response.cancel with response_cancel_not_active when nothing + is answering yet, an error about an event the client never sent.""" + + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + def cancel_not_active(event_id: str) -> bytes: + return json.dumps( + {"type": "error", "error": {"code": "response_cancel_not_active", "event_id": event_id}} + ).encode() + + def proxy_cancel_event_id() -> str: + sent = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + return next(event["event_id"] for event in sent if event["type"] == "response.cancel") + + backend_frames: Final = iter( + ( + lambda: _make_transcript_event("blocked words"), + lambda: cancel_not_active(proxy_cancel_event_id()), + lambda: cancel_not_active("evt_from_client"), + lambda: cancel_not_active("litellm_guardrail_from_client"), + ) + ) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame() + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + client_errors: Final = [ + event["error"] for call in client_ws.send_text.await_args_list if (event := json.loads(call.args[0]))["type"] == "error" + ] + assert [(error["code"], error.get("event_id")) for error in client_errors] == [ + ("content_policy_violation", None), + ("response_cancel_not_active", "evt_from_client"), + ("response_cancel_not_active", "litellm_guardrail_from_client"), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("event_normalizer", "expected_frames"), + [ + (None, ["response.cancel", "conversation.item.delete", "conversation.item.create", "response.create"]), + (XAIRealtimeNormalizer(), ["response.cancel", "conversation.item.create", "response.create"]), + ], + ids=["deletes_the_blocked_item", "xai_keeps_the_item_it_cannot_delete"], +) +async def test_blocked_transcript_is_removed_from_the_conversation_before_the_refusal_turn( + monkeypatch: pytest.MonkeyPatch, + event_normalizer: XAIRealtimeNormalizer | None, + expected_frames: list[str], +): + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_frames: Final = iter((_make_transcript_event("blocked words", item_id="item_blocked"),)) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming( + client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer + ).backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["type"] for event in sent] == expected_frames + assert all(event["item_id"] == "item_blocked" for event in sent if event["type"] == "conversation.item.delete") diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index fa649b48a51..93f700b1c2a 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6143,72 +6143,6 @@ def test_unit_test_perplexity_citations_chunk(): assert "citations" in response_chunk -def test_mock_response_iterator_tool_use(): - """ - Relevant Issue: https://github.com/BerriAI/litellm/issues/7364 - """ - from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator - from litellm.types.utils import ( - ChatCompletionMessageToolCall, - Choices, - CompletionTokensDetailsWrapper, - Function, - Message, - PromptTokensDetailsWrapper, - Usage, - ) - - litellm.set_verbose = False - response = ModelResponse( - id="chatcmpl-Ai8KRI5vJPZXQ9SQvEJfTVuVqkyEZ", - created=1735081811, - model="o1-2024-12-17", - object="chat.completion", - system_fingerprint="fp_e6d02d4a78", - choices=[ - Choices( - finish_reason="tool_calls", - index=0, - message=Message( - content=None, - role="assistant", - tool_calls=[ - ChatCompletionMessageToolCall( - function=Function( - arguments='{"location":"San Francisco, CA","unit":"fahrenheit"}', - name="get_current_weather", - ), - id="call_BfRX2S7YCKL0BtxbWMl89ZNk", - type="function", - ) - ], - function_call=None, - ), - ) - ], - usage=Usage( - completion_tokens=1955, - prompt_tokens=85, - total_tokens=2040, - completion_tokens_details=CompletionTokensDetailsWrapper( - accepted_prediction_tokens=0, - audio_tokens=0, - reasoning_tokens=1920, - rejected_prediction_tokens=0, - text_tokens=None, - ), - prompt_tokens_details=PromptTokensDetailsWrapper( - audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None - ), - ), - service_tier=None, - ) - completion_stream = MockResponseIterator(model_response=response) - response_chunk = completion_stream._chunk_parser(chunk_data=response) - - assert response_chunk["tool_use"] is not None - - def test_is_delta_empty(): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import Delta 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/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 2f0a7ad070b..9e0bd5c60b8 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -287,7 +287,7 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block(): assert block_type == "thinking" assert content_block_start == { "type": "thinking", - "thinking": "I need to summar", + "thinking": "", "signature": "", } diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py new file mode 100644 index 00000000000..f906530c8d9 --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py @@ -0,0 +1,131 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.anthropic.pass_through.adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError +from litellm.router import Router +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +_TOOL_OUTPUT: Final = "line of the file under review " * 40 +_INVALID_COMPACT_EDIT: Final = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _response() -> ModelResponse: + return ModelResponse( + choices=[Choices(index=0, finish_reason="stop", message=Message(content="done"))], + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + + +def _messages() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "review a.py"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "read_file", "input": {"path": "a.py"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": _TOOL_OUTPUT}]}, + ] + + +class _Backend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + async def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +class _SyncBackend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +def _router(deployment_params: Mapping[str, object], router_defaults: Mapping[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **deployment_params}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +async def _send(router: Router, request_options: Mapping[str, object], backend: _Backend) -> object: + with patch("litellm.acompletion", new=backend): + return await router.aanthropic_messages( + model="gemini-agent", + max_tokens=16, + messages=_messages(), + tools=[{"name": "read_file", "input_schema": {"type": "object", "properties": {}}}], + context_management=_INVALID_COMPACT_EDIT, + **request_options, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: Mapping[str, object], + router_defaults: Mapping[str, object], + request_options: Mapping[str, object], + litellm_drop_params: bool, +) -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", litellm_drop_params): + response: Final = await _send(_router(deployment_params, router_defaults), request_options, backend) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) + + +@pytest.mark.asyncio +async def test_messages_rejects_an_invalid_edit_without_drop_params() -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", False), pytest.raises(AnthropicContextManagementError) as error: + await _send(_router({}, {}), {}, backend) + + assert error.value.status_code == 400 + assert backend.calls == [] + + +def test_sync_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit() -> None: + backend: Final = _SyncBackend() + with patch.object(litellm, "drop_params", False), patch("litellm.completion", new=backend): + response: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=16, + messages=_messages(), + model="gemini/gemini-2.5-flash", + api_key="test-key", + context_management=_INVALID_COMPACT_EDIT, + drop_params=True, + ) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py index 9a6940ff12b..b73df255cad 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -12,7 +12,11 @@ content survives. import asyncio import json from types import SimpleNamespace -from typing import AsyncIterator +from typing import AsyncIterator, Final + +import anthropic +import httpx +import pytest from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, @@ -273,3 +277,93 @@ def test_thinking_then_signature_chunk_does_not_crash_stream(): assert "message_stop" in sse assert "Done" in sse + + +def _fake_stream_with_thinking(thinking_blocks: list[dict[str, str]]) -> MockResponseIterator: + model_response: Final = ModelResponse() + model_response.choices = [ + Choices( + index=0, + message=Message( + role="assistant", + content='{"a": 1}', + reasoning_content="".join(block.get("thinking", "") for block in thinking_blocks), + thinking_blocks=thinking_blocks, + ), + finish_reason="stop", + ) + ] + model_response.usage = Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + return MockResponseIterator(model_response=model_response) + + +def _anthropic_sdk_final_content( + completion_stream: "MockResponseIterator | AsyncIterator[ModelResponseStream]", +) -> list[dict[str, object]]: + sse: Final = _collect_async(AnthropicStreamWrapper(completion_stream=completion_stream, model="claude")) + transport: Final = httpx.MockTransport( + lambda request: httpx.Response(200, content=sse.encode(), headers={"content-type": "text/event-stream"}) + ) + client: Final = anthropic.Anthropic(api_key="transport-only", http_client=httpx.Client(transport=transport)) + with client.messages.stream( + model="claude", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) as stream: + return [block.model_dump(exclude_none=True) for block in stream.get_final_message().content] + + +@pytest.mark.parametrize( + "thinking_blocks", + [ + [ + {"type": "thinking", "thinking": "first", "signature": "SIG_A"}, + {"type": "thinking", "thinking": "second", "signature": "SIG_B"}, + ], + [ + {"type": "redacted_thinking", "data": "OPAQUE"}, + {"type": "thinking", "thinking": "visible", "signature": "SIG_C"}, + ], + [ + {"type": "thinking", "thinking": "unsigned", "signature": ""}, + {"type": "thinking", "thinking": "signed", "signature": "SIG_D"}, + ], + ], + ids=["two_signed_thinking_blocks", "redacted_then_signed_thinking", "unsigned_then_signed_thinking"], +) +def test_fake_stream_gives_each_thinking_block_its_own_content_block(thinking_blocks: list[dict[str, str]]): + content: Final = _anthropic_sdk_final_content(_fake_stream_with_thinking(thinking_blocks)) + + assert content == [*thinking_blocks, {"type": "text", "text": '{"a": 1}'}] + + +def _native_thinking_chunk(thinking: str, index: int, signature: str | None = None) -> ModelResponseStream: + block: Final = {"type": "thinking", "thinking": thinking, "index": index, **({"signature": signature} if signature else {})} + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(reasoning_content=thinking or None, thinking_blocks=[block]))] + ) + + +async def _native_stream_with_two_thinking_blocks() -> "AsyncIterator[ModelResponseStream]": + for chunk in ( + _native_thinking_chunk("ab", 0), + _native_thinking_chunk("cd", 0), + _native_thinking_chunk("", 0, "SIG_A"), + _native_thinking_chunk("ef", 1), + _native_thinking_chunk("gh", 1), + _native_thinking_chunk("", 1, "SIG_B"), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ), + ): + yield chunk + + +def test_native_stream_thinking_blocks_keep_their_text_once(): + content: Final = _anthropic_sdk_final_content(_native_stream_with_two_thinking_blocks()) + + assert content == [ + {"type": "thinking", "thinking": "abcd", "signature": "SIG_A"}, + {"type": "thinking", "thinking": "efgh", "signature": "SIG_B"}, + {"type": "text", "text": "hi"}, + ] diff --git a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py index 0a5fbcdf63e..5e55c0ade32 100644 --- a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py @@ -2544,7 +2544,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() the summary subcall lands on the router with no parent attribution, and those tokens go unbilled to the caller's key/team.""" from litellm.llms.anthropic.pass_through.adapters.handler import ( - _prepare_context_managed_request, + prepare_context_managed_request, ) captured_summary_metadata: Dict[str, Any] = {} @@ -2561,7 +2561,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() ), patch("litellm.token_counter", return_value=200_000), ): - result = await _prepare_context_managed_request( + result = await prepare_context_managed_request( model=MODEL, messages=_simple_messages(), tools=None, 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/chat/test_azure_chat_o_series_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py index 3630ef1e354..1cdc288ff87 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -1,9 +1,10 @@ import json import traceback -from typing import Callable, Optional +from typing import Callable, Final, Optional from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import ModelResponse @@ -130,7 +131,7 @@ def test_azure_o3_streaming(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o3-mini", @@ -141,7 +142,7 @@ def test_azure_o3_streaming(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" in mock_create.call_args.kwargs + assert "stream" in mock_create.call_args.kwargs["body"] def test_azure_o_series_routing(): @@ -157,7 +158,7 @@ def test_azure_o_series_routing(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o_series/my-random-deployment-name", @@ -168,17 +169,19 @@ def test_azure_o_series_routing(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" not in mock_create.call_args.kwargs + assert "stream" not in mock_create.call_args.kwargs["body"] @patch("litellm.main.azure_o1_chat_completions._get_openai_client") def test_openai_o_series_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="azure/o1-preview", @@ -233,7 +236,7 @@ async def test_azure_o1_series_response_format_extra_params(): response_format = {"type": "json_object"} tool_choice = "auto" with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -251,7 +254,7 @@ async def test_azure_o1_series_response_format_extra_params(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", json.dumps(request_body, indent=4)) assert request_body["tools"] == tools 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/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7907344ae64..7d5afde68ef 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -313,7 +313,7 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): number: str = Field(description="total number of days in a week") days: list[str] = Field(description="name of days in a week") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: try: completion( model="azure/gpt4o", @@ -326,9 +326,9 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): mock_post.assert_called_once() - print(mock_post.call_args.kwargs) + print(mock_post.call_args.kwargs["body"]) - assert "tool_calls" not in mock_post.call_args.kwargs + assert "tool_calls" not in mock_post.call_args.kwargs["body"] @pytest.mark.parametrize( @@ -382,7 +382,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): } ] - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: mock_post.return_value.headers = {} mock_post.return_value.parse.return_value = litellm.ModelResponse( choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}] @@ -411,9 +411,9 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): mock_post.assert_called_once() if api_version == "2024-10-21": - assert "response_format" in mock_post.call_args.kwargs + assert "response_format" in mock_post.call_args.kwargs["body"] else: - assert "response_format" not in mock_post.call_args.kwargs + assert "response_format" not in mock_post.call_args.kwargs["body"] assert response.choices[0].message.content == InvestigationOutput().model_dump_json() 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..634228da1bf 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,195 @@ 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) +@pytest.mark.parametrize( + ("upstream_model", "usage", "expected_completion_tokens"), + [ + pytest.param( + "grok-4.7", + { + "prompt_tokens": 14, + "completion_tokens": 10, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="grok-reports-reasoning-outside-completion", + ), + pytest.param( + "gpt-5.4-mini", + { + "prompt_tokens": 14, + "completion_tokens": 322, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="reasoning-already-inside-completion", + ), + ], +) +def test_azure_ai_counts_reasoning_tokens_once_whatever_the_deployment_is_named( + respx_mock: respx.MockRouter, + upstream_model: str, + usage: dict[str, object], + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + respx_mock.post(url__startswith=api_url).mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-x", + "object": "chat.completion", + "created": 1700000000, + "model": upstream_model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "42"}, "finish_reason": "stop"}], + "usage": usage, + }, + ) + ) + + response: Final = litellm.completion( + model="azure_ai/team-reasoner", + 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", + ) + + assert response.usage.completion_tokens == expected_completion_tokens + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +@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 + + +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("completion_tokens", "expected_completion_tokens"), + [ + pytest.param(10, 322, id="reasoning-outside-completion"), + pytest.param(322, 322, id="reasoning-already-inside-completion"), + ], +) +def test_azure_ai_stream_counts_reasoning_tokens_once_whatever_the_chunks_call_the_model( + respx_mock: respx.MockRouter, + completion_tokens: int, + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + chunk_base: Final = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "team-reasoner", + } + chunks_sent: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "42"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 14, + "completion_tokens": completion_tokens, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + }, + ) + stream_body: Final = "".join(f"data: {json.dumps({**chunk_base, **chunk})}\n\n" for chunk in chunks_sent) + respx_mock.post(url__startswith=api_url).mock( + return_value=Response(200, text=f"{stream_body}data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + ) + + chunks: Final = tuple( + litellm.completion( + model="azure_ai/team-reasoner", + 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 len(usage_chunks) == 1 + assert usage_chunks[0].completion_tokens == expected_completion_tokens + 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/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index d0219ea2963..924b192a00e 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -1,9 +1,12 @@ +import json from collections.abc import Mapping from typing import Final +import httpx import pytest from pydantic import TypeAdapter, ValidationError +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( ir_to_systemone_request, ir_to_systemone_response, @@ -15,10 +18,13 @@ from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, DecisionsIRRefusal, DecisionsIRRequest, + DecisionsIRScoreAnswer, DecisionsRequestBody, OpenAIDecisionRequestBody, UnsupportedDecisionsRequest, ) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager _SYSTEMONE_BODY: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody) _OPENAI_BODY: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody) @@ -264,3 +270,66 @@ def test_a_box_answer_without_confidence_reaches_the_caller_unchanged() -> None: response: Final = ir_to_systemone_response(parse_systemone_response(payload, ir), ir) assert response.model_dump(mode="json") == payload + + +@pytest.mark.parametrize("provider", ("nace", "openrouter", "fireworks_ai", "typesafe")) +@pytest.mark.parametrize("legend", (["low", "high"], {"0": "low", "1": "high"})) +def test_score_legends_preserve_labels_probabilities_and_metadata(provider: str, legend: object) -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("drex-v2.0", LlmProviders(provider)) + assert config is not None + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate( + { + "state": {"text": "The order arrived damaged"}, + "questions": {"severity": {"type": "score", "criteria": ["low", "high"]}}, + } + ) + ) + payload: Final = { + "model": "drex-v2.0", + "answers": { + "severity": { + "type": "score", + "score": 0.75, + "confidence": 0.8, + "legend": legend, + "probabilities": {"0": 0.25, "1": 0.75}, + } + }, + "usage": {"input_tokens": 20, "output_tokens": 30}, + "request_id": "test-request", + "evaluation_time_ms": 8.5, + } + original: Final = json.dumps(payload) + response: Final = config.parse_response(payload, request) + answer: Final = response.answers[0] + assert isinstance(answer, DecisionsIRScoreAnswer) + assert (answer.score, answer.confidence) == (0.75, 0.8) + assert tuple((item.value, item.label, item.probability) for item in answer.probabilities) == ( + (0, "low", 0.25), + (1, "high", 0.75), + ) + assert response.model == payload["model"] + assert (response.usage.input_tokens, response.usage.output_tokens) == (20, 30) + assert response.extra == {"request_id": "test-request", "evaluation_time_ms": 8.5} + assert json.dumps(payload) == original + + +@pytest.mark.parametrize( + "payload", ([], {"answers": []}, {"answers": {"a": None}}, {"answers": {"a": {"type": "score", "legend": [1]}}}) +) +@pytest.mark.parametrize("provider", ("nace", "openrouter", "fireworks_ai", "typesafe")) +def test_invalid_provider_response_is_an_api_error(payload: object, provider: str) -> None: + config: Final = ProviderConfigManager.get_provider_decisions_config("drex-v2.0", LlmProviders(provider)) + assert config is not None + request: Final = systemone_request_to_ir( + DecisionsRequestBody.model_validate( + { + "state": "text", + "questions": {"a": {"type": "score", "criteria": ["low", "high"]}}, + } + ) + ) + with pytest.raises(BaseLLMException, match="unexpected response") as error: + config.transform_decisions_response("drex-v2.0", provider, httpx.Response(200, json=payload), request) + assert error.value.status_code == 500 diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 936bf06bc18..204b458cadc 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -11,9 +11,9 @@ import httpx import pytest from pydantic import BaseModel -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler -from litellm.types.utils import GenericStreamingChunk +from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse class TestBaseModelResponseIterator: @@ -314,3 +314,25 @@ def test_close_releases_the_attached_http_response(): iterator.http_response = response iterator.close() assert response.is_closed + + +def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_the_response(): + """The OpenAI SDK's `.stream()` accumulator merges streamed `thinking_blocks` entries by `index`, so a fake stream + numbers them the way a real stream does, redacted blocks included""" + thinking_blocks: Final = [ + {"type": "thinking", "thinking": "Step 1.", "signature": "sig-a"}, + {"type": "redacted_thinking", "data": "opaque"}, + {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b"}, + ] + response: Final = ModelResponse( + choices=[Choices(message=Message(content="Paris", thinking_blocks=thinking_blocks), finish_reason="stop")] + ) + + chunk: Final = convert_model_response_to_streaming(response) + + assert chunk.choices[0].delta.thinking_blocks == [ + {"type": "thinking", "thinking": "Step 1.", "signature": "sig-a", "index": 0}, + {"type": "redacted_thinking", "data": "opaque", "index": 1}, + {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2}, + ] + assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None] 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/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index c025a368e51..dec32b3325e 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio from datetime import datetime, timedelta, timezone +from itertools import chain from typing import Final from unittest.mock import MagicMock, patch @@ -302,3 +303,230 @@ async def test_async_converse_streaming_fails_at_the_request_timeout_not_the_ups def test_sync_converse_streaming_fails_at_the_request_timeout_not_the_upstreams_pace() -> None: with pytest.raises(litellm.Timeout): litellm.completion(client=slow_upstream_sync_client(), **_converse_streaming_kwargs()) + + +RECORDED_REASONING: Final = "Tokyo proper has about 14 million people." +RECORDED_THINKING_TOKENS: Final = 309 +JSON_MODE_THINKING_CONVERSE_RESPONSE: Final = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": RECORDED_REASONING, + "signature": "EqoBCkgIBxABGAIiQL1", + } + } + }, + { + "toolUse": { + "toolUseId": "tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", + "name": "json_tool_call", + "input": {"city": "Tokyo", "population_millions": 14.0}, + } + }, + ], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 473, "outputTokens": 390, "totalTokens": 863}, + "additionalModelResponseFields": { + "usage": {"output_tokens_details": {"thinking_tokens": RECORDED_THINKING_TOKENS}} + }, +} + + +def _answer_with_json_mode_thinking_response(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=JSON_MODE_THINKING_CONVERSE_RESPONSE, request=request) + + +def _bedrock_claude_without_native_structured_output() -> str: + return next( + f"bedrock/{model}" + for model, info in litellm.model_cost.items() + if info.get("litellm_provider") == "bedrock_converse" + and model.startswith("us.anthropic.claude") + and info.get("supports_native_structured_output") is False + and info.get("supports_reasoning") is True + ) + + +def _json_mode_thinking_kwargs() -> dict[str, object]: + return { + "model": _bedrock_claude_without_native_structured_output(), + "messages": [{"role": "user", "content": "How many people live in Tokyo, in millions?"}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "city_fact", + "schema": { + "type": "object", + "properties": {"city": {"type": "string"}, "population_millions": {"type": "number"}}, + "required": ["city", "population_millions"], + }, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + } + + +def _answer_fields(response: ModelResponse) -> tuple[object, ...]: + choice: Final = response.choices[0] + return ( + choice.finish_reason, + choice.message.content, + choice.message.tool_calls, + getattr(choice.message, "reasoning_content", None), + getattr(choice.message, "thinking_blocks", None), + getattr(response.usage.completion_tokens_details, "reasoning_tokens", None), + ) + + +def _streamed_thinking_blocks(chunks: list[litellm.ModelResponseStream]) -> list[object]: + deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices] + block_lists: Final = [getattr(delta, "thinking_blocks", None) or [] for delta in deltas] + return list(chain.from_iterable(block_lists)) + + +def test_fake_streamed_json_mode_answer_keeps_the_reasoning_of_the_non_streamed_answer() -> None: + """Without native structured output, json_schema plus streaming fakes the stream from one Converse call. That fake + stream dropped the reasoning, the signed thinking blocks and the reasoning tokens, and ended on tool_calls""" + transport: Final = httpx.MockTransport(_answer_with_json_mode_thinking_response) + client: Final = HTTPHandler(client=httpx.Client(transport=transport)) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list( + litellm.completion(client=client, stream=True, stream_options={"include_usage": True}, **kwargs) + ) + rebuilt: Final = litellm.stream_chunk_builder(chunks, messages=kwargs["messages"]) + + finish_reason, _, _, reasoning_content, _, reasoning_tokens = _answer_fields(answer) + assert (finish_reason, reasoning_content, reasoning_tokens) == ( + "stop", + RECORDED_REASONING, + RECORDED_THINKING_TOKENS, + ) + assert _answer_fields(rebuilt) == _answer_fields(answer) + assert _streamed_thinking_blocks(chunks) == [ + {**block, "index": index} for index, block in enumerate(answer.choices[0].message.thinking_blocks) + ] + + +@pytest.mark.asyncio +async def test_async_fake_streamed_json_mode_answer_keeps_the_reasoning_of_the_non_streamed_answer() -> None: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_answer_with_json_mode_thinking_response)) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = await litellm.acompletion(client=client, **kwargs) + stream: Final = await litellm.acompletion( + client=client, stream=True, stream_options={"include_usage": True}, **kwargs + ) + chunks: Final = [chunk async for chunk in stream] + rebuilt: Final = litellm.stream_chunk_builder(chunks, messages=kwargs["messages"]) + + finish_reason, _, _, reasoning_content, _, reasoning_tokens = _answer_fields(answer) + assert (finish_reason, reasoning_content, reasoning_tokens) == ( + "stop", + RECORDED_REASONING, + RECORDED_THINKING_TOKENS, + ) + assert _answer_fields(rebuilt) == _answer_fields(answer) + assert _streamed_thinking_blocks(chunks) == [ + {**block, "index": index} for index, block in enumerate(answer.choices[0].message.thinking_blocks) + ] + + +LLAMA_TEXT_TOOL_CALL_CONVERSE_RESPONSE: Final = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": '{"type": "function", "name": "get_weather", "parameters": {"city": "Paris"}}'}], + } + }, + "stopReason": "end_turn", + "usage": {"inputTokens": 175, "outputTokens": 23, "totalTokens": 198}, +} + + +def _answer_with_llama_text_tool_call(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=LLAMA_TEXT_TOOL_CALL_CONVERSE_RESPONSE, request=request) + + +def _llama_tool_kwargs() -> dict[str, object]: + return { + "model": "bedrock/us.meta.llama3-3-70b-instruct-v1:0", + "messages": [{"role": "user", "content": "What is the weather in Paris right now?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ], + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + } + + +def _tool_answer(response: ModelResponse) -> tuple[object, ...]: + choice: Final = response.choices[0] + return ( + choice.finish_reason, + choice.message.content, + [(call.function.name, call.function.arguments) for call in choice.message.tool_calls or []], + ) + + +def test_fake_streamed_text_tool_call_arrives_as_the_tool_call_of_the_non_streamed_answer() -> None: + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_answer_with_llama_text_tool_call))) + kwargs: Final = _llama_tool_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list(litellm.completion(client=client, stream=True, **kwargs)) + + assert _tool_answer(answer) == ("tool_calls", None, [("get_weather", '{"city": "Paris"}')]) + assert _tool_answer(litellm.stream_chunk_builder(chunks, messages=kwargs["messages"])) == _tool_answer(answer) + + +@pytest.mark.asyncio +async def test_async_fake_streamed_text_tool_call_arrives_as_the_tool_call_of_the_non_streamed_answer() -> None: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_answer_with_llama_text_tool_call)) + kwargs: Final = _llama_tool_kwargs() + + answer: Final = await litellm.acompletion(client=client, **kwargs) + stream: Final = await litellm.acompletion(client=client, stream=True, **kwargs) + chunks: Final = [chunk async for chunk in stream] + + assert _tool_answer(answer) == ("tool_calls", None, [("get_weather", '{"city": "Paris"}')]) + assert _tool_answer(litellm.stream_chunk_builder(chunks, messages=kwargs["messages"])) == _tool_answer(answer) + + +def test_fake_streamed_json_mode_answer_closes_merged_reasoning_before_the_answer() -> None: + """A fake stream sends reasoning and answer in one chunk, so merging reasoning into content must open and close + the think block around the reasoning in that chunk, before the answer""" + client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(_answer_with_json_mode_thinking_response)) + ) + kwargs: Final = _json_mode_thinking_kwargs() + + answer: Final = litellm.completion(client=client, **kwargs) + chunks: Final = list( + litellm.completion(client=client, stream=True, merge_reasoning_content_in_choices=True, **kwargs) + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == ( + f"{answer.choices[0].message.reasoning_content}{answer.choices[0].message.content}" + ) 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/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 8f95a502137..a181c5a5425 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -17,6 +17,7 @@ from litellm.constants import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.realtime_api import main as realtime_main @pytest.fixture(autouse=True) @@ -109,6 +110,14 @@ class UnavailableBedrockStream: raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") +class RecordingRealtimeHandler: + def __init__(self): + self.calls = 0 + + async def async_realtime(self, **kwargs): + self.calls += 1 + + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id @@ -678,18 +687,23 @@ class TestBedrockRealtimeSessionLifecycle: @pytest.mark.asyncio @pytest.mark.parametrize( - "drop_params,expected_errors", - [(False, [("invalid_type", "session.temperature")]), (True, [])], + "session,drop_params,expected_errors", + [ + ({"temperature": "0.2", "modalities": ["text"]}, False, [("invalid_type", "session.temperature")]), + ({"temperature": "0.2", "modalities": ["text"]}, True, []), + ({"modalities": "text"}, False, [("invalid_type", "session.modalities")]), + ({"modalities": "text"}, True, []), + ], ) - async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( - self, stub_aws_sdk_client, drop_params, expected_errors + async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, session, drop_params, expected_errors ): stream: Final = ScriptedBedrockStream([]) stub_aws_sdk_client["streams"] = [stream] user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} websocket: Final = ConnectedClientWS( [ - json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "session.update", "session": session}), json.dumps({"type": "conversation.item.create", "item": user_text}), ] ) @@ -777,12 +791,26 @@ class TestBedrockRealtimeProviderFailurePropagation: dispatched["logging_obj"] = RecordingLogging() return dispatched + USER_TEXT: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + } + ) + REJECTED_SESSION_UPDATE: Final = json.dumps({"type": "session.update", "session": {"temperature": "0.2"}}) + USER_AUDIO: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_messages", + [[SESSION_UPDATE], [USER_TEXT], [USER_AUDIO], [REJECTED_SESSION_UPDATE, USER_TEXT]], + ids=["session-update-first", "text-first", "audio-first", "input-after-rejected-update"], + ) async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( - self, stub_aws_sdk_client, spend_dispatch + self, stub_aws_sdk_client, spend_dispatch, monkeypatch, client_messages: list[str] ): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.SESSION_UPDATE]) + websocket = ConnectedClientWS(client_messages) stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver) stub_aws_sdk_client["streams"] = [stream] @@ -814,6 +842,37 @@ class TestBedrockRealtimeProviderFailurePropagation: assert not litellm.should_retry(replay.value.status_code), "the router must not retry the replay refusal" assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error" + other_provider: Final = RecordingRealtimeHandler() + monkeypatch.setattr(realtime_main, "openai_realtime", other_provider) + with pytest.raises(BedrockError) as fallback: + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime", websocket=websocket, litellm_logging_obj=FakeLogging() + ) + + assert fallback.value.status_code == 400, "a fallback to another provider must not start a second session" + assert other_provider.calls == 0 + + @pytest.mark.asyncio + async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.USER_TEXT]) + healthy_stream = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + assert failure.value.status_code == 503 + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert any(event["event"].get("textInput", {}).get("content") == "hi" for event in replayed) + @pytest.mark.asyncio async def test_input_side_failure_keeps_spend_for_responses_already_delivered( self, stub_aws_sdk_client, spend_dispatch diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 1cc135a8671..3e10c3bcbdc 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1655,9 +1655,17 @@ class TestBedrockRealtimeUserEventsAndUsage: ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ({"instructions": {"role": "pirate"}}, "session.instructions"), + ({"voice": {"id": "voice_123"}}, "session.voice"), + ({"modalities": "text"}, "session.modalities"), + ({"modalities": ["text", 1]}, "session.modalities"), + ({"input_sample_rate_hertz": "16000"}, "session.input_sample_rate_hertz"), + ({"tools": {"type": "function", "name": "lookup"}}, "session.tools"), + ({"input_audio_format": {"type": "audio/pcm"}}, "session.input_audio_format"), + ({"output_audio_format": {"type": "audio/pcmu"}}, "session.output_audio_format"), ], ) - def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + def test_malformed_setting_is_rejected_without_drop_params(self, session, param): rejection: Final = BedrockRealtimeConfig().session_update_rejection( {"type": "session.update", "session": session} ) @@ -1677,6 +1685,59 @@ class TestBedrockRealtimeUserEventsAndUsage: inference: Final = self._inference_configuration(config, session) assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + def test_non_number_input_sample_rate_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"input_sample_rate_hertz": "8000"} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + config.transform_session_update_event({"type": "session.update", "session": session}) + assert config.input_sample_rate_hertz == BedrockRealtimeConfig().input_sample_rate_hertz + + def test_non_string_instructions_and_voice_are_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"instructions": {"role": "pirate"}, "voice": {"id": "voice_123"}} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + sent: Final = self._voice_and_system_prompt(config, session) + assert sent == self._voice_and_system_prompt(BedrockRealtimeConfig(), {}) + + @pytest.mark.parametrize("voice", ["alloy", "marin", "Alloy", "", " "]) + def test_openai_or_empty_voice_is_rejected_without_drop_params(self, voice): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"voice": voice}} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", "session.voice") + assert f"'{voice}'" in rejection.message + + def test_openai_voice_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"voice": "alloy", "instructions": "Answer as a pirate."} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + voice, system_prompt = self._voice_and_system_prompt(config, session) + assert voice == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})[0] + assert system_prompt == "Answer as a pirate." + + def test_nova_voice_is_applied(self): + session: Final = {"voice": "tiffany"} + assert BedrockRealtimeConfig().session_update_rejection({"type": "session.update", "session": session}) is None + assert self._voice_and_system_prompt(BedrockRealtimeConfig(), session)[0] == "tiffany" + + def test_g711_input_audio_format_sets_the_input_sample_rate(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"input_audio_format": "g711_ulaw"}}) + assert config.input_sample_rate_hertz == 8000 + + @staticmethod + def _voice_and_system_prompt(config: BedrockRealtimeConfig, session: dict[str, object]) -> tuple[object, object]: + events: Final = [ + json.loads(message)["event"] + for message in config.transform_session_update_event({"type": "session.update", "session": session}) + ] + voice: Final = next( + e["promptStart"]["audioOutputConfiguration"]["voiceId"] for e in events if "promptStart" in e + ) + system_prompt: Final = next(e["textInput"]["content"] for e in events if "textInput" in e) + return voice, system_prompt + def test_null_number_settings_count_as_unset(self): session: Final = {"temperature": None, "max_response_output_tokens": None} config: Final = BedrockRealtimeConfig() 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/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py index bcd5f3d8d19..017f279d6c4 100644 --- a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping -from typing import cast +from typing import Final, cast from unittest.mock import MagicMock import pytest @@ -1323,6 +1323,40 @@ def test_gemini_realtime_pipecat_semantic_vad_omits_realtime_input_config(): assert setup["tools"][0]["function_declarations"][0]["name"] == "terminate_call" +@pytest.mark.parametrize( + ("create_response", "expected_detection"), + [ + (False, {"disabled": True}), + (True, {"disabled": False, "prefixPaddingMs": 200, "silenceDurationMs": 700}), + ], +) +def test_gemini_turn_detection_tuning_is_dropped_when_auto_response_is_off( + create_response: bool, expected_detection: dict[str, object] +): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": create_response, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + messages: Final = GeminiRealtimeConfig().transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == expected_detection + + def test_gemini_input_audio_buffer_commit_maps_to_audio_stream_end(): config = GeminiRealtimeConfig() setup = { diff --git a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py index 86fb6fa15d8..9b8788d3769 100644 --- a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py +++ b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py @@ -1,15 +1,16 @@ import json +import re from io import BytesIO from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest from openai import AsyncOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage -import pytest -import re import litellm from litellm import completion @@ -18,7 +19,6 @@ from litellm.llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfi from tests._vcr_conftest_common import rewound_new_episodes_cassette from tests.capturing_transport import CapturingTransport - _GATEWAY_EMBEDDING_RESPONSE: Final = CreateEmbeddingResponse( object="list", data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),), @@ -98,7 +98,7 @@ async def test_litellm_gateway_from_sdk(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -114,9 +114,9 @@ async def test_litellm_gateway_from_sdk(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize("is_async", (False, True)) @@ -264,7 +264,7 @@ async def test_litellm_gateway_from_sdk_rerank(is_async): patch_target = client.post with patch.object(client, "post", new=mock_method): - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) # Create a mock response similar to OpenAI's rerank response mock_response.text = json.dumps( @@ -408,7 +408,7 @@ async def test_litellm_gateway_from_sdk_structured_output(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: litellm.completion( @@ -427,8 +427,8 @@ async def test_litellm_gateway_from_sdk_structured_output(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - json_schema = mock_call.call_args.kwargs["response_format"] + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) + json_schema = mock_call.call_args.kwargs["body"]["response_format"] assert "json_schema" in json_schema @@ -486,7 +486,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): openai_client = OpenAI(api_key="fake-key") # Create mock response object - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) mock_response.headers = {"x-litellm-response-cost": "120"} mock_response.parse.return_value = litellm.ModelResponse( **{ @@ -528,8 +528,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): ) with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", return_value=mock_response, ) as mock_call: response = litellm.completion( @@ -564,7 +563,7 @@ async def test_litellm_gateway_image_generation_direct(is_async): "created": 1, "data": [{"url": "https://example.com/image.png"}], } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.parse.return_value = mock_openai_response mock_raw_response.headers = {} diff --git a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py index a5d7e47fb65..8b46ea53713 100644 --- a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py @@ -20,6 +20,8 @@ from litellm.llms.meta.realtime.transformation import ( session_created_event, ) from litellm.types.llms.meta import MuseMode +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager from litellm.types.realtime import RealtimeResponseTransformInput EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { @@ -164,6 +166,22 @@ def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, ob parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) +def _handshake(turn_detection: str, drop_params: bool) -> tuple[str | bytes, ...]: + config = ProviderConfigManager.get_provider_realtime_config(MUSE_MODEL, LlmProviders.META, drop_params=drop_params) + assert config is not None + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + return tuple(config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MUSE_MODEL)) + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _handshake("semantic_vad", drop_params=True) == _handshake("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(MuseProtocolError, match="Set drop_params to use server_vad instead"): + _handshake("semantic_vad", drop_params=False) + + def test_session_created_event_exposes_openai_transcription_shape(): config = parse_session_update( _event( diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index e5298c7ba39..3633ac97265 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -94,7 +94,7 @@ def test_chat_completion_nvidia_nim_with_tools(): ] with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: completion( @@ -116,7 +116,7 @@ def test_chat_completion_nvidia_nim_with_tools(): # Add assertions to check the request mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -208,7 +208,7 @@ def test_completion_nvidia_nim(): api_key="fake-api-key", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: completion( model=model_name, @@ -226,7 +226,7 @@ def test_completion_nvidia_nim(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["messages"] == [ { 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_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py index 6ee882b7d32..121fab66c15 100644 --- a/tests/unit/llms/openai/test_o_series_transformation.py +++ b/tests/unit/llms/openai/test_o_series_transformation.py @@ -70,7 +70,7 @@ async def test_o1_handle_system_role(model): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -83,7 +83,7 @@ async def test_o1_handle_system_role(model): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -141,7 +141,7 @@ async def test_o1_max_completion_tokens(model: str): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -154,7 +154,7 @@ async def test_o1_max_completion_tokens(model: str): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index e37fd4b6ae5..11d7c7d6d3c 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -5,12 +5,13 @@ import os from collections.abc import Mapping from types import ModuleType from typing import Final, Literal, Protocol -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import httpx import pytest import respx -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AzureOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai._types import Response as SDKResponse from pydantic import BaseModel, ConfigDict, JsonValue @@ -221,6 +222,79 @@ async def test_acompletion_returns_prompt_cache_details_over_injected_transport( assert prompt_tokens_details.cached_tokens == 7 +_REASONING_TOOL_CALL_REPLY: Final = { + "id": "chatcmpl-fake-stream", + "object": "chat.completion", + "created": 0, + "model": "o4-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city": "Paris"}'}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, +} + + +def test_azure_o_series_fake_stream_keeps_every_tool_call_and_token_detail(): + request_bodies: Final[list[dict[str, object]]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + request_bodies.append(json.loads(request.content)) + return httpx.Response(200, json=_REASONING_TOOL_CALL_REPLY) + + client: Final = AzureOpenAI( + api_key="transport-only", + api_version="2025-03-01-preview", + azure_endpoint="https://fake-stream.openai.azure.com", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ) + chunks: Final = list( + litellm.completion( + model="azure/o_series/reasoning-deployment", + api_key="transport-only", + api_version="2025-03-01-preview", + api_base="https://fake-stream.openai.azure.com", + client=client, + messages=[{"role": "user", "content": "weather and time in Paris?"}], + stream=True, + stream_options={"include_usage": True}, + num_retries=0, + ) + ) + rebuilt: Final = litellm.stream_chunk_builder(chunks) + + assert [body.get("stream") for body in request_bodies] == [None] + assert rebuilt is not None + assert rebuilt.choices[0].finish_reason == "tool_calls" + assert [call.function.name for call in rebuilt.choices[0].message.tool_calls] == ["get_weather", "get_time"] + assert rebuilt.usage.prompt_tokens_details.cached_tokens == 64 + assert rebuilt.usage.completion_tokens_details.reasoning_tokens == 1920 + + @respx.mock def test_completion_defaults_missing_message_role_and_preserves_omitted_fields() -> None: route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( @@ -1596,7 +1670,7 @@ async def test_openai_prediction_param_mock(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1615,7 +1689,7 @@ async def test_openai_prediction_param_mock(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the prediction parameter assert "prediction" in request_body @@ -1627,10 +1701,12 @@ async def test_openai_prediction_param_mock(): def test_openai_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="gpt-4o-mini", @@ -1739,7 +1815,7 @@ async def test_openai_safety_identifier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1752,7 +1828,7 @@ async def test_openai_safety_identifier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1768,7 +1844,7 @@ def test_openai_safety_identifier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1781,7 +1857,7 @@ def test_openai_safety_identifier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1798,7 +1874,7 @@ async def test_openai_service_tier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1811,7 +1887,7 @@ async def test_openai_service_tier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1829,7 +1905,7 @@ def test_openai_service_tier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1842,7 +1918,7 @@ def test_openai_service_tier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1874,3 +1950,63 @@ def test_responses_gpt54_with_xhigh_reasoning(): assert request_body["model"] == "openai/gpt-5.4" assert request_body["reasoning"] == {"effort": "xhigh"} + + +_CHAT_RESPONSE_JSON: Final = { + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_WIRE_REQUEST: Final = { + "model": "gpt-5.4", + "messages": [ + {"role": "system", "content": "sys"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is in the image"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png", "detail": "low"}}, + ], + }, + ], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}], + "tool_choice": "auto", + "response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}, + "max_completion_tokens": 7, + "reasoning_effort": "none", + "n": 2, + "seed": 3, +} + + +@pytest.mark.parametrize("is_async", (False, True)) +@respx.mock +async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_headers_unchanged( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_RESPONSE_JSON) + ) + kwargs: Final = { + **_WIRE_REQUEST, + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1", + "api_key": "sk-test", + "num_retries": 0, + "extra_body": {"vendor_flag": True}, + "extra_headers": {"x-caller": "unit"}, + } + + if is_async: + await litellm.acompletion(**kwargs) + else: + litellm.completion(**kwargs) + + sent: Final = route.calls.last.request + assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True} + assert sent.headers["x-caller"] == "unit" + assert sent.headers["authorization"] == "Bearer sk-test" 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_empty_response.py b/tests/unit/llms/openai/test_openai_empty_response.py index 26b28f967db..1505d9ef049 100644 --- a/tests/unit/llms/openai/test_openai_empty_response.py +++ b/tests/unit/llms/openai/test_openai_empty_response.py @@ -2,13 +2,13 @@ Test for issue #17209: Clearer error when LLM endpoint returns empty response """ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest +from openai._legacy_response import LegacyAPIResponse - -from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.openai import OpenAIChatCompletion class TestEmptyResponseHandling: @@ -22,12 +22,12 @@ class TestEmptyResponseHandling: openai_chat = OpenAIChatCompletion() # Mock the raw response to return an empty string from parse() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = "" # Empty string response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -48,12 +48,12 @@ class TestEmptyResponseHandling: """Test that None response also produces a clear error""" openai_chat = OpenAIChatCompletion() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = None mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -75,12 +75,12 @@ class TestEmptyResponseHandling: mock_response = MagicMock() mock_response.model_dump.return_value = {"choices": []} - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -104,12 +104,12 @@ class TestEmptyResponseHandling: # Create a mock response WITHOUT model_dump (like an AsyncStream/Iterator) mock_stream = MagicMock(spec=[]) # spec=[] means no attributes - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_stream mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) 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/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py index 7013afc7a5f..58f1ff60ba5 100644 --- a/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py +++ b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py @@ -9,6 +9,7 @@ body. OpenAI/Anthropic reject unknown body params with HTTP 400. from unittest.mock import MagicMock +from openai._legacy_response import LegacyAPIResponse import litellm from litellm.types.utils import all_litellm_params @@ -47,12 +48,12 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): }, } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -65,7 +66,7 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): ) create_kwargs = ( - mock_client.chat.completions.with_raw_response.create.call_args.kwargs + mock_client.post.call_args.kwargs ) assert "use_chat_completions_api" not in create_kwargs assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {}) diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py new file mode 100644 index 00000000000..75b60b2dd6c --- /dev/null +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -0,0 +1,75 @@ +import json +from typing import Final + +from unittest.mock import Mock + +import httpx + +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openai_like.chat.handler import make_sync_call + +_FAKE_STREAMED_REPLY: Final = { + "id": "chatcmpl-fake-stream", + "object": "chat.completion", + "created": 0, + "model": "openai/gpt-oss-120b", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "reasoning_content": "Look up both values for Paris.", + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }, + { + "id": "call_time", + "type": "function", + "function": {"name": "get_time", "arguments": '{"city": "Paris"}'}, + }, + ], + }, + } + ], + "usage": { + "prompt_tokens": 85, + "completion_tokens": 1955, + "total_tokens": 2040, + "prompt_tokens_details": {"cached_tokens": 64}, + "completion_tokens_details": {"reasoning_tokens": 1920}, + }, +} + + +def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail(): + request_bodies: Final[list[dict[str, object]]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + request_bodies.append(json.loads(request.content)) + return httpx.Response(200, json=_FAKE_STREAMED_REPLY) + + chunks: Final = list( + make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + api_base="https://openai-like.example/v1/chat/completions", + headers={}, + data=json.dumps({"model": "openai/gpt-oss-120b", "messages": []}), + model="openai/gpt-oss-120b", + messages=[], + logging_obj=Mock(), + fake_stream=True, + ) + ) + + assert len(request_bodies) == 1 + assert len(chunks) == 1 + delta: Final = chunks[0].choices[0].delta + assert delta.reasoning_content == "Look up both values for Paris." + assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"] + assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64 + assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920 diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py index 88382b80d16..dbc6ec54a40 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py @@ -1,5 +1,11 @@ +from typing import Final +from unittest.mock import MagicMock + +import httpx import pytest + from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.utils import ModelResponse def test_sanitize_usage_obj_handles_null_tokens(): @@ -47,3 +53,32 @@ def test_sanitize_usage_obj_valid_usage(): # Assert assert sanitized_json == original_json # The object should be unchanged + + +def test_transform_response_surfaces_provider_headers_as_llm_provider_headers(): + raw_response: Final = httpx.Response( + status_code=200, + headers={"x-amzn-requestid": "req_abc123", "content-type": "application/json"}, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + response: Final = OpenAILikeChatConfig().transform_response( + model="google.gemma-4-31b", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req_abc123" diff --git a/tests/unit/llms/perplexity/chat/test_transformation.py b/tests/unit/llms/perplexity/chat/test_transformation.py index 6a70fd1232d..0cb6335f5a8 100644 --- a/tests/unit/llms/perplexity/chat/test_transformation.py +++ b/tests/unit/llms/perplexity/chat/test_transformation.py @@ -4,6 +4,7 @@ import os from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import completion @@ -92,7 +93,7 @@ class TestPerplexityReasoning: pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"content-type": "application/json"} new_response.parse.return_value = pydantic_obj return new_response @@ -100,8 +101,7 @@ class TestPerplexityReasoning: openai_client = OpenAI(api_key="fake-api-key") with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: response = completion( @@ -121,7 +121,7 @@ class TestPerplexityReasoning: # Get the request data from the mock call call_args = mock_client.call_args - request_data = call_args.kwargs + request_data = call_args.kwargs["body"] # Verify reasoning_effort was included in the request assert "reasoning_effort" in request_data diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index 5687a319f06..f9d12b7eeea 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -8,7 +8,7 @@ import os import copy import uuid import json -from typing import Any, Dict, List +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, patch, Mock, MagicMock @@ -595,6 +595,26 @@ class TestSnowflakeCortexClaudeFixes: assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"} assert [b["type"] for b in blocks] == ["thinking", "tool_use"] + def test_signed_thinking_blocks_replay_without_their_stream_index(self): + body: Final = self._transform( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "391", + "thinking_blocks": [ + {"type": "thinking", "thinking": "391", "signature": "Eto", "index": 0}, + {"type": "redacted_thinking", "data": "opaque", "index": 1}, + ], + }, + ] + ) + blocks: Final = body["messages"][1]["content"] + assert blocks[:2] == [ + {"type": "thinking", "thinking": "391", "signature": "Eto"}, + {"type": "redacted_thinking", "data": "opaque"}, + ] + def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self): """A thinking response without a tool call must also round-trip on the next request.""" body = self._transform( diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index b50d6e19458..5c7099832ba 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( parse_chirp_session_update, ) from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config from litellm.types.llms.vertex_ai_speech_to_text import ( VertexSpeechStreamingConfigured, VertexSpeechStreamingResponse, @@ -173,6 +174,26 @@ def test_unsupported_session_settings_are_rejected(payload: str, message: str): parse_chirp_session_update(payload, MODEL) +def _configure_command(turn_detection: str, drop_params: bool) -> list[object]: + config = vertex_realtime_config( + MODEL, access_token="token", resolve_access_token=_token, project="proj-1", location="us", drop_params=drop_params + ) + config.transform_session_created_event(MODEL, "sess_1") + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MODEL) + ] + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _configure_command("semantic_vad", drop_params=True) == _configure_command("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(ChirpProtocolError, match="Set drop_params to use server_vad instead"): + _configure_command("semantic_vad", drop_params=False) + + def test_session_update_configures_once_and_later_updates_are_ignored(): config = _config() config.transform_session_created_event(MODEL, "sess_1") diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 06dacedde67..405b7234c81 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,5 +1,6 @@ import asyncio import base64 +import logging import time import uuid from collections.abc import Generator @@ -18,10 +19,13 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( truncate_tool_name, ) from litellm.llms.vertex_ai.count_tokens.context import ( + CountContextEdits, + InferenceRequestView, VertexGeminiCountContext, VertexGeminiCountContextError, gemini_count_contents, inference_request_view, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -31,6 +35,7 @@ _SCHEMA: Final[JsonValue] = { "properties": {}, "additionalProperties": False, } +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) _FUNCTION_DECLARATION: Final = { "functionDeclarations": [ { @@ -383,8 +388,8 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value( @pytest.mark.parametrize( "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") ) -def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: - view: Final = inference_request_view( +async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = await inference_request_view( ( {"role": "user", "content": "first question " * 50}, {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, @@ -392,16 +397,18 @@ def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(sys ), system, "anthropic", + _NO_EDITS, ) + assert isinstance(view, InferenceRequestView) assert view.messages == [{"role": "user", "content": "latest question"}] assert "Summary of the earlier turns." in str(view.system) assert "Be brief." in str(view.system) -def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: +async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." - view: Final = inference_request_view( + view: Final = await inference_request_view( ( {"role": "user", "content": "What is new?"}, { @@ -426,18 +433,21 @@ def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> Non ), None, "anthropic", + _NO_EDITS, ) + assert isinstance(view, InferenceRequestView) assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) assert [block["type"] for block in assistant_blocks] == ["text"] assert snippet in assistant_blocks[0]["text"] -def test_openai_view_is_the_request_unchanged() -> None: +async def test_openai_view_is_the_request_unchanged() -> None: messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) - view: Final = inference_request_view(messages, "Be brief.", "openai") + view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) + assert isinstance(view, InferenceRequestView) assert view.messages == messages assert view.system == "Be brief." @@ -470,7 +480,7 @@ async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> Non @pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) -def test_anthropic_replayed_web_search_result_is_counted_as_its_text( +async def test_anthropic_replayed_web_search_result_is_counted_as_its_text( custom_llm_provider: Literal["vertex_ai", "gemini"], ) -> None: snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." @@ -498,9 +508,163 @@ def test_anthropic_replayed_web_search_result_is_counted_as_its_text( {"role": "user", "content": "Thanks."}, ) - view: Final = inference_request_view(messages, None, "anthropic") - contents: Final = asyncio.run( - gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + view: Final = await inference_request_view(messages, None, "anthropic", _NO_EDITS) + assert isinstance(view, InferenceRequestView) + contents: Final = await gemini_count_contents( + view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} ) assert snippet in str(contents) + + +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 +_READ_FILE_TOOL: Final[dict[str, object]] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_EDIT: Final[JsonValue] = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, +} +_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} + + +def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: + return ( + *( + turn + for index in range(first_round, first_round + rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ) + + +def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: + return tuple( + _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) + ) + + +def _edits(spec: JsonValue) -> CountContextEdits: + return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) + + +async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: + spec: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] + } + turns: Final = ( + {"role": "user", "content": "first question " * 30000}, + {"role": "assistant", "content": "first answer"}, + *_agent_turns(3), + ) + + with ( + patch( + "litellm.llms.anthropic.pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="gemini/gemini-2.5-flash", + ), + patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, + ): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) + + assert summary.call_count == 0 + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + with ( + patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert any("skipping edits" in record.getMessage() for record in caplog.records) + assert not any(record.levelno >= logging.ERROR for record in caplog.records) + + +async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_before_compact_fails() -> None: + clear_before_compact: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 3}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112"}, + ] + } + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: + low_trigger: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] + } + + view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) + + assert isinstance(view, VertexGeminiCountContextError) + assert view.status_code == 400 + + +async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None: + turns: Final = ( + *_agent_turns(1)[:-1], + { + "role": "user", + "content": [ + { + "type": "search_result", + "source": "https://docs.litellm.ai", + "title": "Docs", + "content": [{"type": "text", "text": "LiteLLM docs"}], + } + ], + }, + ) + + view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic") + assert blocks_error is not None + assert blocks_error.status_code == 400 + assert "search_result" in blocks_error.error_message diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a705c169782..0bfff870d15 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -223,6 +223,36 @@ def test_vertex_session_update_normalizes_ga_remapped_fields( ) +def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off(): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + messages: Final = cfg.transform_realtime_request( + json.dumps(session_update), + _NATIVE_AUDIO_MODEL, + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == { + "disabled": True, + "prefixPaddingMs": 200, + "silenceDurationMs": 700, + } + + # --------------------------------------------------------------------------- # Round-trip test: text-in / text-out via RealTimeStreaming # --------------------------------------------------------------------------- diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index d244fe6bd90..97c0a657ca4 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -28,7 +28,7 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler -from litellm.types.utils import TokenCountMessagesFormat +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat @pytest.mark.asyncio @@ -62,7 +62,9 @@ class _VertexCountTokensPeer: return httpx.Response(200, json={"totalTokens": 37}, request=request) return httpx.Response( self.status_code, - json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + json={ + "error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"} + }, request=request, ) @@ -133,6 +135,60 @@ async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_with assert peer.requests == () +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +def _tool_round(index: int) -> tuple[dict[str, object], ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_applies_the_context_edits_inference_applies() -> None: + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[*_tool_round(0), *_tool_round(1), *_tool_round(2), {"role": "user", "content": "Summarize."}], + contents=None, + deployment={ + "litellm_params": {**_VERTEX_DEPLOYMENT["litellm_params"], "model": "vertex_ai/gemini-2.5-flash"} + }, + request_model="vertex_ai/gemini-2.5-flash", + tools=[ + {"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}} + ], + messages_format="anthropic", + context_management=CountContextManagement( + spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + drop_invalid_edits=False, + ), + ) + + assert result is not None + assert result.total_tokens == 37 + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_TOOL_OUTPUT) == 1 + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location diff --git a/tests/unit/llms/volcengine/test_volcengine.py b/tests/unit/llms/volcengine/test_volcengine.py index d472d52d2db..dff9a0c62b2 100644 --- a/tests/unit/llms/volcengine/test_volcengine.py +++ b/tests/unit/llms/volcengine/test_volcengine.py @@ -2,6 +2,7 @@ import os import sys from unittest.mock import MagicMock, patch +from openai._legacy_response import LegacyAPIResponse from pydantic import BaseModel from litellm.llms.volcengine.chat.transformation import ( @@ -117,7 +118,7 @@ class TestVolcEngineConfig: client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -126,9 +127,7 @@ class TestVolcEngineConfig: } mock_raw_response.parse.return_value = ModelResponse() - with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response - ) as mock_create: + with patch.object(client, "post", MagicMock(return_value=mock_raw_response)) as mock_create: completion( model="volcengine/doubao-seed-1.6", messages=[ @@ -144,11 +143,6 @@ class TestVolcEngineConfig: ) mock_create.assert_called_once() - print(mock_create.call_args.kwargs) + print(mock_create.call_args.kwargs["body"]) # Fixed: thinking disabled should appear in extra_body with original structure - assert ( - "extra_body" in mock_create.call_args.kwargs - and "thinking" in mock_create.call_args.kwargs.get("extra_body", {}) - and mock_create.call_args.kwargs.get("extra_body", {})["thinking"] - == {"type": "disabled"} - ) + assert mock_create.call_args.kwargs["options"]["extra_json"]["thinking"] == {"type": "disabled"} diff --git a/tests/unit/messages/__init__.py b/tests/unit/messages/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c96e759b1a3..c4414003ca5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -4,6 +4,7 @@ import os from collections.abc import Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -12,6 +13,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from litellm.caching.caching import DualCache from litellm.constants import MCP_ALL_TOOLS_WILDCARD from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -23,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, + LiteLLM_UserTable, ProxyException, SpecialHeaders, SpecialMCPServerNames, @@ -11549,3 +11552,117 @@ class TestAccessGroupToolOwnershipAndFaults: patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})), ): assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"] + + +class _EmptyTables: + """Storage edge for a session credential's catalog refresh: the user row the real + ``reload_admitted_user`` reads, and every other table (keys, MCP servers, config) empty, so the + real ``_reload_admitted_key`` fails closed, since a session credential has no key row.""" + + def __init__(self, live_user: LiteLLM_UserTable) -> None: + self.litellm_usertable = SimpleNamespace(find_unique=AsyncMock(return_value=live_user)) + + def __getattr__(self, table: str) -> SimpleNamespace: + empty: Final = SimpleNamespace( + find_unique=AsyncMock(return_value=None), + find_first=AsyncMock(return_value=None), + find_many=AsyncMock(return_value=[]), + count=AsyncMock(return_value=0), + ) + setattr(self, table, empty) + return empty + + +def _session_credential_storage(live_user: LiteLLM_UserTable) -> SimpleNamespace: + tables: Final = _EmptyTables(live_user) + return SimpleNamespace(db=tables, writer_db=tables) + + +@pytest.mark.asyncio +async def test_catalog_refresh_reloads_session_credential_through_its_user_and_keeps_token_grants(monkeypatch): + """A ``/token``-minted or ``lite login`` credential has no key row: the catalog refresh must read the live + user policy, never the key table, and keep the team, routes and models sealed in the token.""" + from litellm.proxy import proxy_server + + database = _session_credential_storage( + LiteLLM_UserTable(user_id="session-user", organization_id="live-org", user_role="internal_user", teams=[]) + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + caller = UserAPIKeyAuth( + api_key="session-token-hash", + user_id="session-user", + team_id="token-team", + org_id="stale-org", + user_role="proxy_admin", + models=["token-model"], + allowed_routes=["mcp_inference_routes"], + is_session_token=True, + ) + caller.via_virtual_key = True + refreshed = await MCPRequestHandler.refresh_catalog_authority(caller) + assert database.writer_db.litellm_usertable.find_unique.await_count >= 1 + database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited() + assert (refreshed.object_permission, refreshed.object_permission_id) == (None, None) + assert (refreshed.user_role, refreshed.org_id) == ("internal_user", "live-org") + assert (refreshed.team_id, refreshed.models, refreshed.allowed_routes) == ( + "token-team", + ["token-model"], + ["mcp_inference_routes"], + ) + assert refreshed.is_session_token and refreshed.requires_fresh_policy + assert caller.org_id == "stale-org" and not caller.requires_fresh_policy + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["litellm:mcp litellm:llm", "litellm:mcp", None]) +async def test_token_minted_credential_lists_tools_for_every_scope_shape(monkeypatch, scope): + """``tools/list`` with the credential ``/token`` mints (a session token the shared auth marks + ``via_virtual_key`` at the key-cache boundary) returns the tool list instead of the JSON-RPC + credential error, for the combined, MCP-only and unscoped ID-JAG shapes.""" + from mcp.types import ListToolsRequest, ListToolsResult, Tool + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.id_jag_grant import SubjectTokenRefusal, granted_scope + from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context + from litellm.proxy._types import MCPTransport + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-session-credential-tools-list") + granted = granted_scope(scope, mcp_resource=False) + assert not isinstance(granted, SubjectTokenRefusal) + user = LiteLLM_UserTable(user_id="okta-user-1", user_email="okta@example.com", user_role="internal_user") + minted = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user, team_id="team-eng", team_alias="eng", allowed_routes=granted.allowed_routes + ) + credential = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(minted) + assert credential is not None and credential.is_session_token + credential.api_key = "hash-of-" + minted[:16] + credential.via_virtual_key = True + server = MCPServer( + server_id="eng-server", + name="eng_server", + server_name="eng_server", + alias="eng_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + ) + database = _session_credential_storage(user) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + tool = Tool(name="eng_server-list_repos", description="list repos", inputSchema={"type": "object"}) + with ( + patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), + patch.object( + operations.global_mcp_server_manager, + "get_tools_page", + AsyncMock(return_value=ListToolsResult(tools=[tool])), + ), + ): + result = await GatewayOperations().execute(ListToolsRequest(), prepare_context(user_api_key_auth=credential)) + assert [listed.name for listed in result.tools] == ["eng_server-list_repos"] + assert database.writer_db.litellm_usertable.find_unique.await_count >= 1 + database.writer_db.litellm_verificationtoken.find_unique.assert_not_awaited() + assert credential.team_id == "team-eng" and credential.allowed_routes == list(granted.allowed_routes or ()) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 2cdf716fa30..9cd19e9a1c9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -31,6 +31,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ClientSecretAuth, CredError, IdJagConfig, + IdJagSubjectSource, NoneConfig, PassthroughConfig, PrivateKeyJwtAuth, @@ -908,3 +909,25 @@ def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_b assert isinstance(spec.config, TokenExchangeConfig) assert spec.config.client_id == "client" assert spec.config.client_secret is None + assert spec.config.client_id == "client" + assert spec.config.client_secret is None + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (None, IdJagSubjectSource.login_assertion), + ("login_assertion", IdJagSubjectSource.login_assertion), + ("refresh_token", IdJagSubjectSource.refresh_token), + ], +) +def test_id_jag_subject_source_maps_to_config(configured: str | None, expected: IdJagSubjectSource) -> None: + spec = to_server_spec(_id_jag_server(id_jag_subject_source=configured)) + + assert spec is not None + assert isinstance(spec.config, IdJagConfig) + assert spec.config.subject_source is expected + + +def test_id_jag_unknown_subject_source_defers_like_a_half_configured_server() -> None: + assert to_server_spec(_id_jag_server(id_jag_subject_source="assertion")) is None diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py new file mode 100644 index 00000000000..ab2b5b3d966 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py @@ -0,0 +1,412 @@ +"""The login-time exchange of a captured identity assertion for an IdP refresh token.""" + +import asyncio +import logging +from datetime import datetime, timedelta, timezone + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion_exchange import ( + exchange_login_assertion_for_refresh_token, + login_exchange_servers, + refresh_token_exchange_targets, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import SSOClientConfig +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SAML2_SUBJECT_TOKEN_TYPE, + SSOIdentityAssertion, + assertion_from_saml_login, + login_refresh_token_expired, + with_login_refresh_token, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ExchangedToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + ClientSecretAuth, + CredError, + IdJagConfig, + IdJagResourceClient, + IdJagSubjectSource, + NoneConfig, + ServerSpec, +) +from litellm.proxy._types import LiteLLM_MCPServerTable +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_NOW = datetime(2026, 10, 16, 12, 0, tzinfo=timezone.utc) + + +def _config( + *, source: IdJagSubjectSource = IdJagSubjectSource.refresh_token, client_id: str = "agent", endpoint: str = "" +) -> IdJagConfig: + return IdJagConfig( + org_token_endpoint=endpoint or "https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id=client_id, + client_auth=ClientSecretAuth(client_secret=SecretStr("s")), + resource_client=IdJagResourceClient( + client_id=client_id, client_auth=ClientSecretAuth(client_secret=SecretStr("s")) + ), + subject_source=source, + ) + + +def _spec(config: IdJagConfig | NoneConfig, server_id: str = "srv") -> ServerSpec: + return ServerSpec(server_id=server_id, resource="https://mcp.example.com", config=config) + + +def _saml() -> SSOIdentityAssertion: + return assertion_from_saml_login(b"", int(_NOW.timestamp()) + 300, "https://idp.example.com") + + +class _Endpoint: + def __init__(self, result: Result[ExchangedToken, CredError]) -> None: + self._result = result + self.calls: list[tuple[str, str, dict[str, str], ClientAuth]] = [] + + async def fetch( + self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: + self.calls.append((endpoint, client_id, dict(grant_params), client_auth)) + return self._result + + +def test_targets_keep_only_refresh_token_sourced_id_jag_servers_one_per_idp_client(): + direct = _config(source=IdJagSubjectSource.login_assertion) + first = _config(client_id="agent") + twin = _config(client_id="agent") + other = _config(client_id="other-agent") + + targets = refresh_token_exchange_targets( + [_spec(direct, "a"), None, _spec(NoneConfig(), "b"), _spec(first, "c"), _spec(other, "d"), _spec(twin, "e")] + ) + + assert targets == (first, other) + + +@pytest.mark.asyncio +async def test_exchange_presents_the_captured_assertion_and_keeps_the_refresh_token(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + config = _config() + saml = _saml() + + exchanged = await exchange_login_assertion_for_refresh_token(saml, (config,), endpoint, _NOW) + + ((url, client_id, params, auth),) = endpoint.calls + assert (url, client_id, auth) == (config.org_token_endpoint, config.client_id, config.client_auth) + assert params == { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token", + "subject_token": saml.id_token.get_secret_value(), + "subject_token_type": SAML2_SUBJECT_TOKEN_TYPE, + "scope": "openid offline_access", + } + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + assert exchanged.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert exchanged.exchanged_refresh_token_client_id == "agent" + assert exchanged.exchanged_refresh_token_expires_at == _NOW + timedelta(seconds=3600) + assert exchanged.id_token == saml.id_token + assert exchanged.expires_at == saml.expires_at + assert exchanged.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert not login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3599)) + assert login_refresh_token_expired(exchanged, _NOW + timedelta(seconds=3600)) + + +@pytest.mark.asyncio +async def test_exchange_without_a_lifetime_leaves_expiry_to_the_idp(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=None))) + + exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (_config(),), endpoint, _NOW) + + assert exchanged.exchanged_refresh_token_expires_at is None + assert not login_refresh_token_expired(exchanged, _NOW + timedelta(days=365)) + + +@pytest.mark.asyncio +async def test_a_refused_exchange_keeps_the_plain_assertion(caplog: pytest.LogCaptureFixture): + endpoint = _Endpoint(Error(CredError.of_upstream_unavailable("token exchange failed with status 400"))) + saml = _saml() + + with caplog.at_level("WARNING"): + exchanged = await exchange_login_assertion_for_refresh_token(saml, (_config(),), endpoint, _NOW) + + assert exchanged == saml + assert exchanged.exchanged_refresh_token is None + assert any("refused to exchange the login assertion" in message for message in caplog.messages) + + +@pytest.mark.asyncio +async def test_no_exchange_without_a_refresh_token_sourced_server(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + saml = _saml() + + assert await exchange_login_assertion_for_refresh_token(saml, (), endpoint, _NOW) == saml + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_an_oidc_login_is_exchanged_at_the_id_jag_client_and_keeps_its_sso_refresh_token(): + """The SSO refresh token belongs to the SSO client and cannot be redeemed by the ID-JAG client, so + an OIDC login still exchanges its id_token and both tokens live side by side.""" + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + exchanged = await exchange_login_assertion_for_refresh_token(oidc, (_config(),), endpoint, _NOW) + + assert endpoint.calls[0][2]["subject_token"] == "id-token" + assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:id_token" + assert exchanged.refresh_token is not None and exchanged.refresh_token.get_secret_value() == "sso-rt" + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +def _sso_client(client_id: str = "agent", endpoint: str = "https://idp.example.com/token") -> SSOClientConfig: + return SSOClientConfig( + token_endpoint=endpoint, client_id=client_id, client_secret=SecretStr("s"), auth_method="client_secret_basic" + ) + + +@pytest.mark.asyncio +async def test_an_oidc_login_through_the_id_jag_client_keeps_its_own_refresh_token_without_an_exchange(): + """Okta refuses requested_token_type=refresh_token for an id_token subject. When the SSO app is the + ID-JAG client itself, the refresh token the login returned is already that client's, so it is bound + to the target and nothing is posted to the IdP.""" + endpoint = _Endpoint( + Error( + CredError.of_upstream_unavailable( + "token exchange failed with status 400: 'requested_token_type' is invalid or not supported." + ) + ) + ) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + kept = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + assert kept.exchanged_refresh_token_expires_at is None + assert kept.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.id_token == oidc.id_token and kept.refresh_token == oidc.refresh_token + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sso_client", + [ + pytest.param(_sso_client(client_id="sso-app"), id="other-client"), + pytest.param(_sso_client(endpoint="https://other-idp.example.com/token"), id="other-endpoint"), + pytest.param(None, id="unknown-client"), + ], +) +async def test_a_login_refresh_token_from_another_client_or_issuer_is_never_presented( + sso_client: SSOClientConfig | None, +): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + exchanged = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=sso_client + ) + + assert [(call[0], call[1], call[2]["subject_token"]) for call in endpoint.calls] == [ + ("https://idp.example.com/token", "agent", "id-token") + ] + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +@pytest.mark.asyncio +async def test_the_login_refresh_token_binds_to_the_target_whose_client_issued_it_not_the_first(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + targets = (_config(client_id="other-agent"), _config(client_id="agent")) + + kept = await exchange_login_assertion_for_refresh_token(oidc, targets, endpoint, _NOW, sso_client=_sso_client()) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + + +@pytest.mark.asyncio +async def test_a_saml_login_through_the_id_jag_client_still_exchanges_its_assertion(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + + exchanged = await exchange_login_assertion_for_refresh_token( + _saml(), (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:saml2" + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +@pytest.mark.asyncio +async def test_a_login_already_exchanged_is_not_exchanged_again(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-2", expires_in=3600))) + done = with_login_refresh_token(_saml(), "rt-1", 3600, _NOW) + + assert await exchange_login_assertion_for_refresh_token(done, (_config(),), endpoint, _NOW) == done + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_several_idp_clients_mint_for_the_first_and_warn(caplog: pytest.LogCaptureFixture): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + first, second = _config(client_id="agent"), _config(client_id="other-agent") + + with caplog.at_level("WARNING"): + exchanged = await exchange_login_assertion_for_refresh_token(_saml(), (first, second), endpoint, _NOW) + + assert [call[1] for call in endpoint.calls] == ["agent"] + assert exchanged.exchanged_refresh_token is not None + assert any("2 IdP clients" in message and "client_id=agent" in message for message in caplog.messages) + + +class _FakeCatalog: + def __init__(self, config_servers: dict[str, MCPServer]) -> None: + self.config_mcp_servers = config_servers + self.built: list[tuple[str, bool]] = [] + + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: + self.built.append((mcp_server.server_id, register_oauth_discovery)) + return MCPServer( + server_id=mcp_server.server_id, name=mcp_server.server_id, url=mcp_server.url, transport=MCPTransport.http + ) + + +def _row(server_id: str, auth_type: MCPAuth, approval_status: str | None = None) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, + alias=server_id, + url=f"https://{server_id}.example.com/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + approval_status=approval_status, + created_at=_NOW, + updated_at=_NOW, + ) + + +@pytest.mark.asyncio +async def test_login_exchange_servers_reads_config_and_db_id_jag_rows_not_the_registry(): + """A server added on another pod is in the DB before this pod's registry reloads, so the login + exchange has to judge against the DB rows, like the retention gate, and only the ID-JAG ones.""" + config_server = MCPServer( + server_id="from-config", name="from-config", url="https://cfg.example.com/mcp", transport=MCPTransport.http + ) + catalog = _FakeCatalog({"from-config": config_server}) + + servers = await login_exchange_servers( + catalog, + [ + _row("db-id-jag", MCPAuth.oauth2_id_jag), + _row("db-id-jag-approved", MCPAuth.oauth2_id_jag, "approved"), + _row("db-oauth", MCPAuth.oauth2), + ], + ) + + assert [server.server_id for server in servers] == ["from-config", "db-id-jag", "db-id-jag-approved"] + assert catalog.built == [("db-id-jag", False), ("db-id-jag-approved", False)] + + +@pytest.mark.asyncio +async def test_a_db_row_hides_the_config_server_with_the_same_id_like_the_live_catalog(): + """The catalog serves ``config_mcp_servers | registry``, so a DB row wins its server_id. The login + exchange must mint for that live server, not for the hidden config entry's IdP client.""" + hidden = MCPServer( + server_id="shared", name="shared", url="https://cfg.example.com/mcp", transport=MCPTransport.http + ) + catalog = _FakeCatalog({"shared": hidden}) + + servers = await login_exchange_servers(catalog, [_row("shared", MCPAuth.oauth2_id_jag)]) + + assert [(server.server_id, server.url) for server in servers] == [("shared", "https://shared.example.com/mcp")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("approval_status", ["pending_review", "rejected", "draft"]) +async def test_login_exchange_servers_skip_unapproved_submissions(approval_status: str): + """A team key can submit an ID-JAG server naming any token endpoint. Until an admin approves it, + no login assertion is posted there.""" + catalog = _FakeCatalog({}) + + servers = await login_exchange_servers(catalog, [_row("submitted", MCPAuth.oauth2_id_jag, approval_status)]) + + assert servers == () + assert catalog.built == [] + + +class _StalledEndpoint(_Endpoint): + """An IdP that accepts the connection and never answers.""" + + async def fetch( + self, endpoint: str, client_id: str, grant_params: dict[str, str], client_auth: ClientAuth + ) -> Result[ExchangedToken, CredError]: + await asyncio.sleep(0.5) + return await super().fetch(endpoint, client_id, grant_params, client_auth) + + +@pytest.mark.asyncio +async def test_a_stalled_idp_does_not_hold_up_the_login(caplog: pytest.LogCaptureFixture): + """The exchange is optional. A user who only wants the dashboard must not wait on an MCP IdP that + stopped answering, so the exchange has its own short deadline and the login keeps the plain assertion.""" + saml = _saml() + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + exchanged = await exchange_login_assertion_for_refresh_token( + saml, + (_config(),), + _StalledEndpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))), + _NOW, + deadline_seconds=0.05, + ) + + assert exchanged == saml + assert any("did not answer the login-time refresh token exchange" in r.getMessage() for r in caplog.records) + + +class _OneBadRowCatalog(_FakeCatalog): + async def build_mcp_server_from_table( + self, + mcp_server: LiteLLM_MCPServerTable, + *, + credentials_are_encrypted: bool = True, + env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, + ) -> MCPServer: + if mcp_server.server_id == "broken": + raise ValueError("credentials cannot be decrypted") + return await super().build_mcp_server_from_table( + mcp_server, + credentials_are_encrypted=credentials_are_encrypted, + env_vars_are_encrypted=env_vars_are_encrypted, + register_oauth_discovery=register_oauth_discovery, + ) + + +@pytest.mark.asyncio +async def test_one_unbuildable_db_row_does_not_cost_the_other_servers_their_exchange(caplog: pytest.LogCaptureFixture): + catalog = _OneBadRowCatalog({}) + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + servers = await login_exchange_servers( + catalog, [_row("broken", MCPAuth.oauth2_id_jag), _row("healthy", MCPAuth.oauth2_id_jag)] + ) + + assert [server.server_id for server in servers] == ["healthy"] + assert any("server_id=broken" in r.getMessage() for r in caplog.records) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 283d0754940..527475b86a9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -32,6 +32,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, IdJagConfig, IdJagResourceClient, + IdJagSubjectSource, NoneConfig, NoOpAuth, Ok, @@ -55,8 +56,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r SSOClientConfig, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + SAML2_SUBJECT_TOKEN_TYPE, AssertionStoreUnavailable, SSOIdentityAssertion, + assertion_from_saml_login, + with_login_refresh_token, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -111,6 +115,10 @@ class _FakeAssertionStore: self._assertions = dict(assertions or {}) self.lookups: list[str] = [] + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: self.lookups.append(user_id) return self._assertions.get(user_id) @@ -539,12 +547,8 @@ async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): assert leg1_client_id == "litellm" assert isinstance(endpoint.auths[0], ClientSecretAuth) assert endpoint.auths[0].client_secret.get_secret_value() == "s" - assert ( - leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" - ) - assert ( - leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" - ) + assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" assert leg1_params["subject_token"] == "user-id-token" leg2_endpoint, leg2_client_id, leg2_params = endpoint.calls[1] @@ -563,9 +567,7 @@ async def test_id_jag_without_inbound_token_or_stored_assertion_is_precondition_ endpoint = _FakeTokenEndpoint([]) store = _FakeAssertionStore() provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -581,9 +583,7 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) assert _emitted(result.ok)["Authorization"] == "Bearer final-access" @@ -593,6 +593,54 @@ async def test_id_jag_exchanges_the_stored_sso_assertion_when_the_caller_present assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" +@pytest.mark.asyncio +async def test_id_jag_declares_a_stored_saml_assertion_as_a_saml2_subject_token(): + """A SAML login retains the signed assertion, not an id_token, so leg 1 has to tell the IdP what + it is being handed even though the server is configured with the id_token default.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + saml = assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com") + store = _FakeAssertionStore({"alice": saml}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == saml.id_token.get_secret_value() + assert leg1_params["subject_token_type"] == SAML2_SUBJECT_TOKEN_TYPE + assert leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_sends_the_configured_type_for_a_stored_oidc_assertion(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token")}) + config = _id_jag_config().model_copy(update={"subject_token_type": "urn:ietf:params:oauth:token-type:jwt"}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(config)) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" + + +@pytest.mark.asyncio +async def test_id_jag_refuses_a_saml_assertion_past_its_not_on_or_after_without_calling_the_idp(): + endpoint = _FakeTokenEndpoint([]) + saml = assertion_from_saml_login(b"", int(time.time()) - 1, "https://idp.example.com") + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": saml}) + ) + + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + @pytest.mark.asyncio async def test_id_jag_prefers_the_callers_own_token_over_the_stored_assertion(): endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) @@ -613,9 +661,7 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id store = _FakeAssertionStore({"alice": _assertion("stale-id-token", expires_in=-timedelta(seconds=1))}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -646,6 +692,9 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging( rows[user_id] = assertion class _Inner: + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: return await _read(user_id) @@ -672,9 +721,7 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging( ), ) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) _, _, leg1_params = endpoint.calls[0] @@ -686,7 +733,7 @@ def test_the_resolver_defaults_to_the_renewing_assertion_store(): the plain database reader would strand every agent an id_token lifetime after its user's login.""" provider = UpstreamCredentialProvider() - assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # noqa: SLF001 # the wiring is the assertion + assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # the wiring is the assertion @pytest.mark.asyncio @@ -695,9 +742,7 @@ async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): store = _FakeAssertionStore({"alice": _assertion("undated-id-token", expires_in=None)}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Ok) _, _, leg1_params = endpoint.calls[0] @@ -730,14 +775,10 @@ async def test_id_jag_keeps_store_sourced_bearers_partitioned_per_user(): Ok(ExchangedToken(access_token="bob-bearer", expires_in=3600)), ] ) - store = _FakeAssertionStore( - {"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")} - ) + store = _FakeAssertionStore({"alice": _assertion("alice-id-token"), "bob": _assertion("bob-id-token")}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - alice = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + alice = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) bob = await provider.resolve_credentials(Subject(tenant_id="", subject_id="bob"), _spec(_id_jag_config())) assert isinstance(alice, Ok) and isinstance(bob, Ok) @@ -754,6 +795,10 @@ class _OutageAssertionStore: def __init__(self) -> None: self.lookups: list[str] = [] + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: self.lookups.append(user_id) raise AssertionStoreUnavailable(_DRIVER_DETAIL) @@ -766,9 +811,7 @@ async def test_id_jag_maps_an_assertion_store_outage_to_upstream_unavailable(): endpoint = _FakeTokenEndpoint([]) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=_OutageAssertionStore()) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -784,9 +827,7 @@ async def test_id_jag_store_outage_does_not_leak_driver_detail_to_the_caller(cap ) with caplog.at_level(logging.WARNING): - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert _DRIVER_DETAIL not in result.error.summary @@ -803,9 +844,7 @@ async def test_id_jag_invalidation_survives_an_assertion_store_outage(): token_endpoint=_FakeTokenEndpoint([]), sso_assertion_store=_OutageAssertionStore() ) - await provider.invalidate_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) class _FlakyAssertionStore: @@ -815,6 +854,10 @@ class _FlakyAssertionStore: self._assertion = assertion self.down = False + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: if self.down: raise AssertionStoreUnavailable("connection refused") @@ -855,6 +898,10 @@ class _SwitchableAssertionStore: def __init__(self, id_token: str) -> None: self.id_token = id_token + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: return _assertion(self.id_token) @@ -899,6 +946,10 @@ class _SequentialAssertionStore: self.issued: list[str] = [] self._n = 0 + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + + return await self.fetch(user_id) + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: await asyncio.sleep(0) if self.pinned is not None: @@ -996,14 +1047,10 @@ async def test_id_jag_reads_a_naive_stored_expiry_as_utc(): comparison would be a TypeError on the egress path, turning a 412 into a 500.""" endpoint = _FakeTokenEndpoint([]) naive_past = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(hours=1) - store = _FakeAssertionStore( - {"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)} - ) + store = _FakeAssertionStore({"alice": SSOIdentityAssertion(id_token=SecretStr("stale"), expires_at=naive_past)}) provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) - result = await provider.resolve_credentials( - Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "precondition_required" @@ -1032,13 +1079,9 @@ async def test_invalidate_evicts_a_store_sourced_id_jag_bearer(): @pytest.mark.asyncio async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): - endpoint = _FakeTokenEndpoint( - [Error(CredError.of_upstream_unavailable("leg1 down"))] - ) + endpoint = _FakeTokenEndpoint([Error(CredError.of_upstream_unavailable("leg1 down"))]) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -1055,9 +1098,7 @@ async def test_id_jag_propagates_a_leg2_error(): ] ) provider = UpstreamCredentialProvider(token_endpoint=endpoint) - result = await provider.resolve_credentials( - _with_inbound("user-id-token"), _spec(_id_jag_config()) - ) + result = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) assert isinstance(result, Error) assert result.error.tag == "upstream_unavailable" @@ -1065,7 +1106,7 @@ async def test_id_jag_propagates_a_leg2_error(): assert len(endpoint.calls) == 2 -def _two_leg_ok(bearer: str) -> list: +def _two_leg_ok(bearer: str) -> list[Result[ExchangedToken, CredError]]: return [ Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), Ok(ExchangedToken(access_token=bearer, expires_in=3600)), @@ -1256,8 +1297,18 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): (NoneConfig(), _SUBJECT, "no-auth", None), (PassthroughConfig(), _SUBJECT, "no-auth", None), (PassthroughConfig(), _with_inbound("Bearer caller-token"), "oauth2-passthrough", "Bearer caller-token"), - (ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), _SUBJECT, "static-token", "Bearer static-key"), - (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), + ( + ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-key"))), + _SUBJECT, + "static-token", + "Bearer static-key", + ), + ( + AuthorizationCodeConfig(), + Subject(tenant_id="", subject_id="alice"), + "stored-user-token", + "Bearer stored-alice", + ), ], ) async def test_resolved_source_matches_the_credential_sent_upstream( @@ -1279,7 +1330,9 @@ async def test_resolved_source_matches_the_credential_sent_upstream( async def test_resolved_source_preserves_missing_user_token_error() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source - result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) + result = await resolve_credentials_with_source( + UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig()) + ) assert isinstance(result, Error) assert result.error.tag == "unauthorized" @@ -1289,7 +1342,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) - m2m = await resolve_credentials_with_source(UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M)) + m2m = await resolve_credentials_with_source( + UpstreamCredentialProvider(client_credentials_source=source), _SUBJECT, _spec(_M2M) + ) assert isinstance(m2m, Ok) headers, _ = await _emitted_async(m2m.ok.auth) assert headers["Authorization"] == "Bearer m2m-at" @@ -1297,7 +1352,9 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> Non assert source.gets == ["s"] exchanger = _FakeExchanger(Ok(OAuthToken(access_token="exchanged-at"))) - exchanged = await resolve_credentials_with_source(UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO)) + exchanged = await resolve_credentials_with_source( + UpstreamCredentialProvider(token_exchanger=exchanger), _with_inbound("subject"), _spec(_OBO) + ) assert isinstance(exchanged, Ok) assert _emitted(exchanged.ok.auth)["Authorization"] == "Bearer exchanged-at" assert exchanged.ok.source.value == "token-exchange" @@ -1321,9 +1378,14 @@ async def test_id_jag_source_describes_final_token_after_both_exchanges() -> Non async def test_client_credentials_subscription_identity_preserves_resolved_grant_scopes(): async def resolve(token): source = _FakeM2MSource(Ok(token)) - result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials(_SUBJECT, _spec(_M2M)) + result = await UpstreamCredentialProvider(client_credentials_source=source).resolve_credentials( + _SUBJECT, _spec(_M2M) + ) assert isinstance(result, Ok) - from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ClientCredentialsBearerAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + assert isinstance(result.ok, ClientCredentialsBearerAuth) return result.ok.subscription_identity @@ -1333,3 +1395,175 @@ async def test_client_credentials_subscription_identity_preserves_resolved_grant assert original is not None assert original == renewed assert original != reduced + + +def _login_with_refresh_token(expires_in: int | None) -> SSOIdentityAssertion: + saml = assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com") + return with_login_refresh_token( + saml, + "rt-from-login", + expires_in, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + + +def _refresh_token_source_config() -> IdJagConfig: + return _id_jag_config().model_copy(update={"subject_source": IdJagSubjectSource.refresh_token}) + + +@pytest.mark.asyncio +async def test_id_jag_presents_the_login_refresh_token_when_the_server_asks_for_it(): + """``id_jag_subject_source: refresh_token`` presents the refresh token the login assertion was + exchanged for, typed as one, never the SAML assertion itself.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + store = _FakeAssertionStore({"alice": _login_with_refresh_token(3600)}) + provider = UpstreamCredentialProvider(token_endpoint=endpoint, sso_assertion_store=store) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == "rt-from-login" + assert leg1_params["subject_token_type"] == "urn:ietf:params:oauth:token-type:refresh_token" + assert leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_never_presents_the_sso_refresh_token(): + """An OIDC login holds a refresh token for the SSO client. Only the token the ID-JAG client + exchanged at login is presented, the SSO one cannot be redeemed there.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + row = with_login_refresh_token( + oidc, + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_refuses_a_token_minted_for_another_idp_client(): + """A refresh token is bound to the client that received it. One minted for another server's IdP + client is never posted to this server's endpoint, the user is told to sign in again instead.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + row = with_login_refresh_token( + _login_with_refresh_token(3600), + "rt-for-other", + 3600, + datetime.now(timezone.utc), + endpoint="https://other-idp.example.com/token", + client_id="other-agent", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert "different IdP client" in str(result.error) + assert endpoint.calls == [] + + +class _RenewalFailsStore(_FakeAssertionStore): + """An SSO renewal that is timing out: ``fetch`` raises the way RefreshingSSOAssertionStore does, + while the row itself is still readable.""" + + async def fetch_stored(self, user_id: str) -> SSOIdentityAssertion | None: + return self._assertions.get(user_id) + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + raise AssertionStoreUnavailable("IdP token endpoint timed out") + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_does_not_wait_on_an_sso_renewal(): + """The exchanged refresh token has its own lifetime. A failing renewal of the unrelated SSO + id_token must not turn into a 503 for a server that never presents that id_token.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_RenewalFailsStore({"alice": _login_with_refresh_token(3600)}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +async def test_id_jag_refresh_token_source_outlives_the_saml_assertion(): + """The SAML assertion is spent minutes after login; the refresh token it bought is what counts.""" + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + expired_saml = assertion_from_saml_login(b"", int(time.time()) - 60, "https://idp.example.com") + row = with_login_refresh_token( + expired_saml, + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="litellm", + ) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Ok) + assert endpoint.calls[0][2]["subject_token"] == "rt-from-login" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("row", "expected"), + [ + ( + assertion_from_saml_login(b"", int(time.time()) + 300, "https://idp.example.com"), + "none is stored", + ), + (_login_with_refresh_token(-1), "has expired"), + ], + ids=["no refresh token", "expired refresh token"], +) +async def test_id_jag_refresh_token_source_is_a_precondition_without_a_live_refresh_token( + row: SSOIdentityAssertion, expected: str +): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, sso_assertion_store=_FakeAssertionStore({"alice": row}) + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_refresh_token_source_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert expected in result.error.summary + assert "Sign in through LiteLLM SSO" in result.error.summary + assert endpoint.calls == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py index d80913f8d33..53d0d82a7f2 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -35,8 +35,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_r sso_client_config, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + with_login_refresh_token, AssertionStoreUnavailable, SSOIdentityAssertion, + assertion_from_saml_login, ) SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx" @@ -318,6 +320,53 @@ async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(c assert "offline_access" in caplog.text +@pytest.mark.asyncio +async def test_an_expired_saml_assertion_is_never_redeemed_and_names_the_re_login(caplog: pytest.LogCaptureFixture): + """SAML has no refresh grant. Even a row that somehow carries a refresh token must not be posted + to the generic OIDC client, and the operator hint about GENERIC_SCOPE does not apply.""" + saml = assertion_from_saml_login(b"", int(time.time()) - 1, "https://idp.example.com") + with_refresh_token = saml.model_copy(update={"refresh_token": SecretStr("rt_saml")}) + rows = _FakeRows({"alice": with_refresh_token}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch("alice") + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=rows.write) + with caplog.at_level(logging.WARNING): + outcome = await refresher.refresh("alice", with_refresh_token) + + assert served == with_refresh_token + assert transport.calls == [] + assert rows.writes == [] + match outcome: + case Error(failure): + assert failure.kind == "rejected" + assert "sign in" in failure.detail + case Ok(_): + pytest.fail("a SAML assertion must not be renewable") + assert "sign in" in caplog.text + assert "GENERIC_SCOPE" not in caplog.text + + +class _ForbiddenCoordinator: + async def run(self, user_id: str, key: str, *, refresh: object, reread: object) -> None: + raise AssertionError(f"a SAML assertion must never enter the refresh path (user_id={user_id}, key={key})") + + +@pytest.mark.asyncio +async def test_a_saml_assertion_inside_the_expiry_skew_is_served_as_is_without_electing_a_renewal(): + """SAML cannot be renewed, so a still-valid assertion near its NotOnOrAfter is handed out untouched + instead of running the single-flight election that would 503 the losers.""" + saml = assertion_from_saml_login(b"", int(time.time()) + 30, "https://idp.example.com") + rows = _FakeRows({"alice": saml}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport, coordinator_factory=_ForbiddenCoordinator).fetch("alice") + + assert served == saml + assert transport.calls == [] + assert rows.writes == [] + + @pytest.mark.asyncio async def test_an_unconfigured_sso_client_never_calls_the_idp(caplog): rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) @@ -694,6 +743,9 @@ async def test_the_in_process_coordinator_is_retried_until_redis_appears(): {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"}, {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"}, {"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + {"GOOGLE_CLIENT_ID": "google-cid"}, + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "s"}, + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_TENANT": "tenant"}, ], ) def test_a_partial_sso_client_is_no_client(env): @@ -717,6 +769,64 @@ def test_the_configured_sso_client_is_the_one_the_login_used(): assert config.client_secret.get_secret_value() == "s3cret" +def test_a_google_login_is_renewed_at_googles_token_endpoint_as_the_google_client(): + """fastapi-sso's Google login redeems the code at Google's token endpoint with HTTP Basic, so + that is where and how the refresh token it handed back is redeemed.""" + config = sso_client_config({"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"}) + + assert config is not None + assert config.token_endpoint == "https://oauth2.googleapis.com/token" + assert config.client_id == "google-cid" + assert config.client_secret.get_secret_value() == "google-secret" + assert config.auth_method == "client_secret_basic" + + +def test_a_microsoft_login_is_renewed_at_the_tenant_token_endpoint_as_the_microsoft_client(): + config = sso_client_config( + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant-123"} + ) + + assert config is not None + assert config.token_endpoint == "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token" + assert config.client_id == "ms-cid" + assert config.client_secret.get_secret_value() == "ms-secret" + assert config.auth_method == "client_secret_basic" + + +def test_a_custom_microsoft_token_endpoint_wins_over_the_tenant_template(): + """``MICROSOFT_TOKEN_ENDPOINT`` is what the Microsoft login itself honours for sovereign + clouds, so the renewal has to go to the same place.""" + config = sso_client_config( + { + "MICROSOFT_CLIENT_ID": "ms-cid", + "MICROSOFT_CLIENT_SECRET": "ms-secret", + "MICROSOFT_TENANT": "tenant-123", + "MICROSOFT_TOKEN_ENDPOINT": "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token", + } + ) + + assert config is not None + assert config.token_endpoint == "https://login.microsoftonline.us/tenant-123/oauth2/v2.0/token" + + +def test_the_client_follows_the_login_callbacks_provider_order(): + """A deployment carrying a Google and a generic client signs users in through Google, so the + stored refresh token belongs to the Google client and only that client can redeem it.""" + config = sso_client_config( + { + "GOOGLE_CLIENT_ID": "google-cid", + "GOOGLE_CLIENT_SECRET": "google-secret", + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + } + ) + + assert config is not None + assert config.client_id == "google-cid" + assert config.token_endpoint == "https://oauth2.googleapis.com/token" + + def test_the_live_store_renews_over_the_database_reader(): """The composition root has to produce a renewing store, or none of this runs in production.""" assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore) @@ -792,3 +902,43 @@ async def test_a_successful_grant_is_handed_back_as_the_parsed_body(): assert isinstance(outcome, Ok) assert outcome.ok["id_token"] == "idt" + + +@pytest.mark.asyncio +async def test_a_renewal_keeps_the_refresh_token_exchanged_at_login_in_the_stored_row(): + """Renewal replaces the id_token and the SSO refresh token. The refresh token the ID-JAG client + exchanged at login is unrelated to either and must still be in the row the next MCP call reads.""" + stale, fresh = _id_token(exp_offset=-1), _id_token() + row = with_login_refresh_token( + _stored(stale, expires_in=-1), + "rt-from-login", + 3600, + datetime.now(timezone.utc), + endpoint="https://idp.example.com/token", + client_id="agent", + ) + rows = _FakeRows({"alice": row}) + + served = await _store(rows, _FakeTransport(_renewal(fresh, "rt_2"))).fetch("alice") + + stored = rows.rows["alice"] + assert stored.id_token.get_secret_value() == fresh + assert stored.refresh_token is not None and stored.refresh_token.get_secret_value() == "rt_2" + assert stored.exchanged_refresh_token is not None + assert stored.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert stored.exchanged_refresh_token_client_id == "agent" + assert served is not None and served.exchanged_refresh_token is not None + + +@pytest.mark.asyncio +async def test_fetch_stored_hands_back_the_row_without_renewing_it(): + """The read a refresh-token server uses: the row as stored, no IdP round trip even when the SSO + id_token is about to expire.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch_stored("alice") + + assert served is not None + assert served.id_token.get_secret_value() == rows.rows["alice"].id_token.get_secret_value() + assert transport.calls == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index 088b232f3b6..3c8e7a9b448 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -8,9 +8,11 @@ rotation re-encrypts stored rows like the sibling per-user credential tables. """ import asyncio +import base64 import json import os import time +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, patch import jwt as pyjwt @@ -18,15 +20,20 @@ import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( _ASSERTION_CACHE, + SAML2_SUBJECT_TOKEN_TYPE, AssertionStoreUnavailable, DbSSOAssertionStore, SSOAssertionCache, + assertion_from_saml_login, assertion_from_sso_login, + carry_exchanged_refresh_token, ema_assertion_retention_enabled, fetch_sso_identity_assertion, + login_refresh_token_expired, persist_sso_identity_assertion, retain_sso_identity_assertion_for_ema, rotate_sso_identity_assertions_master_key, + with_login_refresh_token, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper from litellm.types.mcp import MCPAuth @@ -111,6 +118,41 @@ def test_assertion_from_sso_login_happy_path(): assert assertion.expires_at.timestamp() == pytest.approx(time.time() + 3600, abs=5) +SAML_ASSERTION_XML = ( + b'' + b"https://idp.example.com/metadata" +) + + +def _saml_subject_bytes(encoded: str) -> bytes: + assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40] + return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + + +def test_assertion_from_saml_login_carries_the_base64url_assertion_as_a_saml2_subject_token(): + not_on_or_after = int(time.time()) + 300 + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, not_on_or_after, "https://idp.example.com/metadata") + assert _saml_subject_bytes(assertion.id_token.get_secret_value()) == SAML_ASSERTION_XML + assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert assertion.refresh_token is None + assert assertion.issuer == "https://idp.example.com/metadata" + assert assertion.expires_at is not None + assert assertion.expires_at.tzinfo is not None + assert assertion.expires_at.timestamp() == not_on_or_after + + +def test_assertion_from_saml_login_without_not_on_or_after_has_no_expiry(): + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, None, None) + assert assertion.expires_at is None + assert assertion.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + + +def test_oidc_login_assertion_leaves_the_subject_token_type_to_the_operator(): + assertion = assertion_from_sso_login(_make_id_token(), None) + assert assertion is not None + assert assertion.subject_token_type is None + + def test_assertion_repr_never_leaks_token_material(): token = _make_id_token() assertion = assertion_from_sso_login(token, "rt_secret_value") @@ -226,6 +268,37 @@ async def test_persist_and_fetch_round_trip_encrypted_at_rest(): assert json.loads(decrypted)["id_token"] == token +@pytest.mark.asyncio +async def test_saml_assertion_round_trips_with_its_subject_token_type(): + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + assertion = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com") + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-saml", assertion) + fetched = await fetch_sso_identity_assertion("user-saml") + assert fetched is not None + assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert _saml_subject_bytes(fetched.id_token.get_secret_value()) == SAML_ASSERTION_XML + assert fetched.expires_at == assertion.expires_at + assert fetched.refresh_token is None + assert SAML_ASSERTION_XML.decode() not in stored["user-saml"] + + +@pytest.mark.asyncio +async def test_oidc_row_written_before_subject_token_types_reads_back_without_one(): + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + token = _make_id_token() + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-a", assertion_from_sso_login(token, None)) + fetched = await fetch_sso_identity_assertion("user-a") + assert fetched is not None + assert fetched.subject_token_type is None + decrypted = decrypt_value_helper(stored["user-a"], "test", exception_type="debug") + assert decrypted is not None + assert "subject_token_type" not in json.loads(decrypted) + + @pytest.mark.asyncio async def test_persist_overwrites_previous_login(): stored = {} @@ -248,7 +321,9 @@ async def test_fetch_serves_second_read_from_cache_without_db_read(): cache = SSOAssertionCache() token = _make_id_token() assertion = assertion_from_sso_login(token, "rt_1") - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion, cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) second = await fetch_sso_identity_assertion("user-a", cache=cache) @@ -266,7 +341,9 @@ async def test_persist_busts_cache_so_relogin_is_visible_immediately(): cache = SSOAssertionCache() first_token = _make_id_token(exp_offset=100) second_token = _make_id_token(exp_offset=7200) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache) @@ -295,7 +372,9 @@ async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion(): return row prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache)) await db_read_started.wait() @@ -324,7 +403,9 @@ async def test_fetch_missing_row_returns_none(): async def test_fetch_does_not_cache_a_missing_row(): prisma = _make_prisma({}) cache = SSOAssertionCache() - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary first = await fetch_sso_identity_assertion("nobody", cache=cache) second = await fetch_sso_identity_assertion("nobody", cache=cache) assert first is None @@ -353,7 +434,9 @@ async def test_cached_assertion_expires_after_ttl(): cache = SSOAssertionCache(ttl_seconds=1) first_token = _make_id_token(exp_offset=100) second_token = _make_id_token(exp_offset=7200) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache) first = await fetch_sso_identity_assertion("user-a", cache=cache) await persist_sso_identity_assertion( @@ -477,7 +560,9 @@ async def test_db_store_uses_injected_cache(): prisma = _make_prisma(stored) cache = SSOAssertionCache() token = _make_id_token() - with patch("litellm.proxy.proxy_server.prisma_client", prisma): # test-quality-ok: fake DB seam has no injection boundary + with patch( + "litellm.proxy.proxy_server.prisma_client", prisma + ): # test-quality-ok: fake DB seam has no injection boundary await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache) store = DbSSOAssertionStore(cache=cache) first = await store.fetch("alice") @@ -496,3 +581,79 @@ async def test_db_store_returns_none_for_a_user_with_no_stored_assertion(): with patch.dict(os.environ, {"LITELLM_SALT_KEY": SALT_KEY}): with patch("litellm.proxy.proxy_server.prisma_client", _make_prisma({})): assert await DbSSOAssertionStore().fetch("nobody") is None + + +@pytest.mark.asyncio +async def test_login_refresh_token_round_trips_beside_the_saml_assertion(): + """A SAML login exchanged for a refresh token keeps the assertion, its type and its expiry, and adds + the refresh token with the lifetime the IdP gave it, all encrypted at rest.""" + stored: dict[str, str] = {} + prisma = _make_prisma(stored) + now = datetime.now(timezone.utc).replace(microsecond=0) + saml = assertion_from_saml_login(SAML_ASSERTION_XML, int(time.time()) + 300, "https://idp.example.com") + assertion = with_login_refresh_token( + saml, "rt-from-login", 3600, now, endpoint="https://idp.example.com/token", client_id="agent" + ) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await persist_sso_identity_assertion("user-rt", assertion) + fetched = await fetch_sso_identity_assertion("user-rt") + assert fetched is not None + assert fetched.exchanged_refresh_token is not None + assert fetched.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert fetched.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600) + assert fetched.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert fetched.exchanged_refresh_token_client_id == "agent" + assert fetched.refresh_token is None + assert fetched.id_token == saml.id_token + assert fetched.expires_at == saml.expires_at + assert fetched.subject_token_type == SAML2_SUBJECT_TOKEN_TYPE + assert "rt-from-login" not in stored["user-rt"] + assert not login_refresh_token_expired(fetched, now) + assert login_refresh_token_expired(fetched, now + timedelta(seconds=3600)) + + +def test_renewal_follows_a_rotated_login_refresh_token_that_is_the_id_jag_subject(): + """When the OIDC login's own refresh token is the ID-JAG subject (the SSO client is the ID-JAG + client), a renewal that rotated it leaves the old token revoked at the IdP, so the subject moves + to the token the renewal returned.""" + now = datetime.now(timezone.utc).replace(microsecond=0) + login = assertion_from_sso_login(_make_id_token(), "sso-rt") + assert login is not None + previous = with_login_refresh_token( + login, "sso-rt", None, now, endpoint="https://idp.example.com/token", client_id="agent" + ) + renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2") + assert renewed is not None + + carried = carry_exchanged_refresh_token(previous, renewed) + + assert carried.exchanged_refresh_token is not None + assert carried.exchanged_refresh_token.get_secret_value() == "sso-rt-2" + assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert carried.exchanged_refresh_token_client_id == "agent" + + +def test_renewal_carries_the_exchanged_refresh_token_forward(): + """An OIDC renewal rebuilds the row from the new id_token. The refresh token the ID-JAG client + exchanged at login is unrelated to that renewal and must survive it.""" + now = datetime.now(timezone.utc).replace(microsecond=0) + previous = with_login_refresh_token( + assertion_from_sso_login(_make_id_token(), "sso-rt"), + "rt-from-login", + 3600, + now, + endpoint="https://idp.example.com/token", + client_id="agent", + ) + renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2") + assert previous is not None and renewed is not None + + carried = carry_exchanged_refresh_token(previous, renewed) + + assert carried.id_token == renewed.id_token + assert carried.refresh_token is not None and carried.refresh_token.get_secret_value() == "sso-rt-2" + assert carried.exchanged_refresh_token is not None + assert carried.exchanged_refresh_token.get_secret_value() == "rt-from-login" + assert carried.exchanged_refresh_token_expires_at == now + timedelta(seconds=3600) + assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert carried.exchanged_refresh_token_client_id == "agent" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f09409a023d..5fdc7669496 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -102,9 +102,8 @@ def isolate_global_mcp_registry(monkeypatch): restore, entries leak into other test modules sharing the same worker and break assertions over the full registry contents. """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - from litellm.proxy._experimental.mcp_server.catalog import CatalogSnapshots + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager monkeypatch.setattr(global_mcp_server_manager, "catalog", CatalogSnapshots(global_mcp_server_manager)) snapshot = dict(global_mcp_server_manager.registry) @@ -8133,9 +8132,9 @@ async def test_resolve_active_litellm_key_resolves_key_without_user_id(proxy_glo blocked and expiry, and the key hash (not the user) is what the mint seals. The per-user token store still gets no user for such a key, since there is none to key a stored credential by.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - extract_user_id_from_request, - _ResolvedKey, _resolve_active_litellm_key, + _ResolvedKey, + extract_user_id_from_request, ) from litellm.proxy._types import UserAPIKeyAuth, hash_token from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -8417,8 +8416,8 @@ async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_rea cached row answers without a database read, and only a caller that asks for the database row pays for one.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - reload_active_user_by_id, load_active_user_by_id, + reload_active_user_by_id, ) from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -11428,11 +11427,10 @@ async def test_upstream_resource_sent_on_dcr_bridge_relay_authorize(): """The DCR-bridge relay arm builds its own upstream authorize params, so it needs the resource too. Without it the relayed authorize would omit the resource while the gateway's token leg still sent one, which is itself an invalid_target.""" - from litellm.types.mcp import MCPAuth - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _dcr_bridge_relays_client_registration, ) + from litellm.types.mcp import MCPAuth server = _resource_server( client_id=None, @@ -11693,7 +11691,7 @@ def _native_client_app(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.master_key", _NATIVE_CLIENT_MASTER_KEY, raising=False) minted = [] - async def fake_mint(user_id, team_id): + async def fake_mint(user_id, team_id, allowed_routes=None): minted.append((user_id, team_id)) return MintedProxyCredential(key=f"sk-cli-{len(minted)}", expires_in=3600, user_id=user_id, team_id=team_id) @@ -12136,7 +12134,9 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch): async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch): from http.cookies import SimpleCookie from urllib.parse import parse_qs, urlparse + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _oauth_state_cookie_name, authorize_with_server, @@ -12211,6 +12211,7 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal @pytest.mark.asyncio async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch): from types import SimpleNamespace + from litellm.caching.caching import DualCache from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import db, mcp_server_manager @@ -12279,9 +12280,10 @@ async def test_enforced_login_warms_verified_token_readable_without_database_loo async def test_dcr_refusal_is_actionable_without_upstream_body( upstream_status: int, auth_type: MCPAuth, dcr_bridge: bool, flow: str, monkeypatch: pytest.MonkeyPatch ) -> None: - import httpx from typing import Final + import httpx + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( mint_ephemeral_dcr_client, register_client_with_server, @@ -12522,8 +12524,8 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user( import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - from litellm.proxy import proxy_server from litellm.models.user import LiteLLM_UserTable + from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.types.llms.custom_http import httpxSpecialProvider @@ -12618,8 +12620,8 @@ async def test_oauth_jwt_identity_rejects_untrusted_or_inactive_owner( from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( - extract_user_id_from_request, authorize_oauth_credential_request, + extract_user_id_from_request, ) allowed_servers: Final = AsyncMock(return_value=["server-a"]) @@ -12688,7 +12690,7 @@ async def test_oauth_jwt_uses_configured_virtual_key_owner( ) -> None: from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.bridge_token_flow import extract_user_id_from_request - from litellm.proxy._types import UserAPIKeyAuth, UnregisteredJWTClientBehavior, hash_token + from litellm.proxy._types import UnregisteredJWTClientBehavior, UserAPIKeyAuth, hash_token from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key handler, signing_key = jwt_oauth_identity @@ -13069,8 +13071,8 @@ async def test_signed_oauth_callback_honors_credential_write_policy( cimd: bool, ) -> None: import httpx - import litellm + import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server @@ -13186,17 +13188,16 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions( allowed: bool, credential: str, ) -> None: - import jwt from datetime import datetime, timedelta, timezone from urllib.parse import parse_qs, urlparse + import jwt + from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server - from litellm.proxy._types import JWTIssuerConfig, UserAPIKeyAuth, hash_token - from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken - from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager - from litellm.proxy._types import MCPTransport + from litellm.proxy._types import JWTIssuerConfig, MCPTransport, UserAPIKeyAuth, hash_token + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer handler, signing_key = jwt_oauth_identity @@ -13330,9 +13331,10 @@ async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session( credential: str, cookie_state: str, ) -> None: - import jwt from urllib.parse import parse_qs, urlparse + import jwt + from litellm.models.user import LiteLLM_UserTable from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager @@ -13664,6 +13666,7 @@ def test_persisted_client_is_not_applied_to_another_upstream(issuer, url): async def test_registration_does_not_write_client_after_upstream_edit(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "test-mcp-registration-salt") from prisma import models + from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _persist_dcr_client_registration @@ -13750,6 +13753,7 @@ async def test_callback_does_not_forward_error_from_another_issuer(monkeypatch): @pytest.mark.parametrize("operation", ["authorize", "token"]) async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, operation): from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -13796,6 +13800,7 @@ async def test_bound_client_cannot_be_sent_to_a_different_issuer(monkeypatch, op @pytest.mark.asyncio async def test_failed_registration_preserves_cached_credentials(monkeypatch): from fastapi import Request + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints server = _dcr_redirect_test_server("old-client").model_copy( @@ -13819,8 +13824,13 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): request.headers = {} with pytest.raises(HTTPException) as error: await endpoints.register_client_with_server( - request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, + request=request, + mcp_server=server, + client_name="app", + grant_types=["authorization_code"], + response_types=["code"], + token_endpoint_auth_method="none", + replace_stale_client=True, ) assert error.value.status_code == 503 assert "could not be saved" in error.value.detail @@ -13835,6 +13845,7 @@ async def test_failed_registration_preserves_cached_credentials(monkeypatch): @pytest.mark.asyncio async def test_registration_without_database_keeps_client_in_temporary_server(monkeypatch): from fastapi import Request + from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints @@ -13850,8 +13861,13 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo request.base_url = "https://gateway.example/" request.headers = {} response = await endpoints.register_client_with_server( - request=request, mcp_server=server, client_name="app", grant_types=["authorization_code"], - response_types=["code"], token_endpoint_auth_method="none", replace_stale_client=True, + request=request, + mcp_server=server, + client_name="app", + grant_types=["authorization_code"], + response_types=["code"], + token_endpoint_auth_method="none", + replace_stale_client=True, ) assert response.status_code == 200 assert json.loads(response.body)["client_id"] == "temporary-client" @@ -13864,9 +13880,10 @@ async def test_registration_without_database_keeps_client_in_temporary_server(mo @pytest.mark.asyncio @pytest.mark.parametrize("config_store", [False, True]) async def test_registration_does_not_overwrite_credentials_after_failed_identity_read(monkeypatch, config_store): - from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db - from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager from litellm.proxy import utils + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager server = _dcr_redirect_test_server(None) monkeypatch.setattr(utils, "get_prisma_client_or_throw", lambda _: MagicMock()) @@ -13895,9 +13912,10 @@ async def test_registration_does_not_overwrite_credentials_after_failed_identity @pytest.mark.asyncio @pytest.mark.parametrize("winner_available", [False, True]) async def test_registration_losing_conditional_write_reuses_only_a_matching_winner(monkeypatch, winner_available): - from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints, db - from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.proxy import utils + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy._types import LiteLLM_MCPServerTable server = _dcr_redirect_test_server(None).model_copy(update={"url": "https://upstream.example/mcp"}) row = LiteLLM_MCPServerTable.model_validate(server.model_dump(exclude_none=True)) @@ -14503,7 +14521,8 @@ async def test_managed_authorize_retains_the_callers_previous_grant( from urllib.parse import parse_qs, urlsplit from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14548,7 +14567,8 @@ async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted from urllib.parse import parse_qs, urlsplit from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14692,7 +14712,8 @@ async def test_managed_refresh_retains_only_the_actual_grant( from types import SimpleNamespace from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager _, signing_key = jwt_oauth_identity @@ -14774,7 +14795,8 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch ], ) async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): - from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints credential = {**_stored_grant(), "scopes": granted} read = AsyncMock(return_value=credential) @@ -14836,3 +14858,137 @@ def test_aggregate_protected_resource_advertises_header_server_scopes( ) with patch.object(mcp_server_manager, "global_mcp_server_manager", manager): assert _build_aggregate_protected_resource_response(request)["scopes_supported"] == expected + + +@pytest.mark.asyncio +async def test_token_endpoint_routes_the_jwt_bearer_grant_to_the_aggregate_gateway(): + """A requesting app presents its identity provider's client_id, never a gateway DCR one, so + the aggregate dispatch keys on the grant type for RFC 7523 and hands the assertion to the + ID-JAG redeemer. A per-server token endpoint never serves that grant.""" + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.id_jag_grant import IdJagPrerequisites + + request = _mock_callback_request("https://litellm.example.com/") + with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id="0oa-requesting-app", + client_secret="app-secret", + assertion="eyJ.assertion.sig", + resource=None, + ) + assert response.status_code == 400 + refused = json.loads(bytes(response.body)) + assert refused["error"] == "unsupported_grant_type" + unmet = IdJagPrerequisites( + jwt_auth_enabled=False, + has_database=False, + licensed=False, + maps_jwts_to_virtual_keys=False, + shared_replay_store=False, + base_url=None, + requesting_apps=(), + ).refusal() + assert unmet is not None and refused["error_description"] == unmet.description + with pytest.raises(HTTPException): + await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id="0oa-requesting-app", + assertion="eyJ.assertion.sig", + mcp_server_name="no-such-server", + ) + + +@pytest.mark.asyncio +async def test_token_endpoint_takes_the_client_id_from_http_basic_when_the_form_leaves_it_out(): + """RFC 6749 section 2.3.1: a confidential app may authenticate with HTTP Basic alone, so the + aggregate dispatch reads client_id from that header. With neither the form field nor the + header the request is a 400 invalid_request, not a 422.""" + from base64 import b64encode + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + + request = _mock_callback_request("https://litellm.example.com/") + request.headers = {"authorization": "Basic " + b64encode(b"0oa-requesting-app:app-secret").decode()} + with patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit9670"): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type=JWT_BEARER_GRANT_TYPE, + client_id=None, + assertion="eyJ.assertion.sig", + resource=None, + ) + assert response.status_code == 400 + assert json.loads(bytes(response.body))["error"] == "unsupported_grant_type" + request.headers = {} + with pytest.raises(HTTPException) as refused: + await discoverable_endpoints.token_endpoint( + request=request, grant_type=JWT_BEARER_GRANT_TYPE, client_id=None, assertion="eyJ.assertion.sig" + ) + assert refused.value.status_code == 400 + assert refused.value.detail["error"] == "invalid_request" + + +@pytest.mark.parametrize("available", [True, False]) +def test_root_metadata_is_the_gateways_own_and_lists_the_jwt_bearer_grant_when_it_is_served(available): + """With no OAuth2 MCP server resolving at the root, the bare-origin document is the gateway's + own authorization server metadata: issuer is the origin an ID-JAG names as ``aud``, the token + endpoint is ``{base}/token``, and the jwt-bearer grant (with HTTP Basic client auth) is listed + exactly when the gateway serves it, under the same gates as the aggregate ``/mcp`` document.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + global_mcp_server_manager.registry.clear() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + with ( + patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=available), # test-quality-ok: deployment gate seam + patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam + ): + document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None) + assert document["issuer"] == "https://llm.example.com" + assert document["token_endpoint"] == "https://llm.example.com/token" + assert (JWT_BEARER_GRANT_TYPE in document["grant_types_supported"]) is available + assert ("client_secret_basic" in document["token_endpoint_auth_methods_supported"]) is available + assert ("none" in document["token_endpoint_auth_methods_supported"]) is available + assert "urn:ietf:params:oauth:grant-type:token-exchange" in document["grant_types_supported"] + + +def test_root_metadata_owned_by_a_resolved_oauth2_server_never_lists_the_gateways_grant(): + """When exactly one OAuth2 MCP server resolves at the root, the root document describes that + upstream's token endpoint, which cannot redeem an ID-JAG, so the jwt-bearer grant is not + advertised there even though the gateway serves it; a requesting app uses the aggregate + ``/.well-known/oauth-authorization-server/mcp`` document in that deployment.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import JWT_BEARER_GRANT_TYPE + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_oauth2_server() + global_mcp_server_manager.registry.clear() + global_mcp_server_manager.registry[server.server_id] = server + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + try: + with ( + patch.object(discoverable_endpoints, "jwt_bearer_available", return_value=True), # test-quality-ok: deployment gate seam + patch.object(discoverable_endpoints, "token_exchange_available", return_value=True), # test-quality-ok: deployment gate seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=server), # test-quality-ok: resolver seam + ): + document = discoverable_endpoints._build_oauth_authorization_server_response(request=request, mcp_server_name=None) + finally: + global_mcp_server_manager.registry.clear() + assert document["token_endpoint"] == "https://llm.example.com/test_oauth/token" + assert JWT_BEARER_GRANT_TYPE not in document["grant_types_supported"] + assert document["grant_types_supported"] == ["authorization_code", "refresh_token"] + assert document["token_endpoint_auth_methods_supported"] == ["client_secret_post"] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1417e77c8e2..ef95df23ff8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -20,10 +20,12 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( CONNECT_FLOW_COOKIE_PREFIX, GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, + JWT_BEARER_GRANT_TYPE, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, SUBJECT_TOKEN_TYPES, TOKEN_EXCHANGE_GRANT_TYPE, + AssertionIdentity, ConsentTeam, MintedProxyCredential, SubjectIdentity, @@ -44,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( open_gateway_dcr_client, register_aggregate_client, revoke_refresh_token, + supported_grant_types, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( SessionBearerAdmitted, @@ -1771,10 +1774,12 @@ CONSENT_TEAMS = (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam class _Minter: def __init__(self, result=None): self.calls = [] + self.allowed_routes = [] self.result = result - async def __call__(self, user_id, team_id): + async def __call__(self, user_id, team_id, allowed_routes=None): self.calls.append((user_id, team_id)) + self.allowed_routes.append(allowed_routes) if self.result is not None: return self.result return MintedProxyCredential(key=f"sk-cli-{user_id}", expires_in=3600, user_id=user_id, team_id=team_id) @@ -3296,3 +3301,200 @@ async def test_gateway_scope_still_narrows_when_the_resource_does_not_resolve() _, cookies = _flow_cookie_from(response) sealed_flow: Final = _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") assert sealed_flow["resource_server_id"] == "alpha-id" + + +ID_JAG = "eyJhbGciOiJSUzI1NiIsInR5cCI6Im9hdXRoLWlkLWphZytqd3QifQ.eyJzdWIiOiJ1MSJ9.idp-signature" +REQUESTING_APP = "0oa-requesting-app" +ASSERTED = AssertionIdentity( + user_id="u1", + team_id="team-b", + replay_key="mcp_gateway_dcr:id_jag:jti-1", + replay_ttl_seconds=300, + scope="litellm:mcp", + allowed_routes=("mcp_inference_routes",), +) + + +class _Redeemer: + def __init__(self, result=None): + self.calls = [] + self.result = result + + async def __call__(self, assertion, client_id, client_secret, request): + self.calls.append((assertion, client_id, client_secret, request.url.path)) + return self.result if self.result is not None else ASSERTED + + +async def _redeem_assertion_native(minter, redeemer, cache=None, client_id=REQUESTING_APP, **overrides): + arguments = { + "grant_type": JWT_BEARER_GRANT_TYPE, + "assertion": ID_JAG, + "client_secret": "app-secret", + "redeem_assertion": redeemer, + } + return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_mints_the_proxy_credential_for_the_asserted_user_once_per_assertion(): + """RFC 7523: an allowlisted requesting app redeems the ID-JAG its identity provider minted for + the same credential the consent flow mints, attributed to the user the assertion names and + the team asked for, with the assertion's scope stated and no refresh token (RFC 7521 section + 4.1, the app renews with a fresh assertion). The assertion's jti is single-use: presenting it again is a replay and mints + nothing, while a fresh assertion for the same user is served.""" + minter, redeemer, cache = _Minter(), _Redeemer(), DualCache() + response = await _redeem_assertion_native(minter, redeemer, cache=cache) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + body = json.loads(response.body) + assert redeemer.calls == [(ID_JAG, REQUESTING_APP, "app-secret", "/token")] + assert minter.calls == [("u1", "team-b")] + assert body["access_token"] == "sk-cli-u1" + assert body["token_type"] == "Bearer" + assert body["scope"] == "litellm:mcp" + assert "issued_token_type" not in body + assert (body["user_id"], body["team_id"]) == ("u1", "team-b") + assert "refresh_token" not in body + replayed = await _redeem_assertion_native(minter, redeemer, cache=cache) + assert replayed.status_code == 400 + assert json.loads(replayed.body)["error"] == "invalid_grant" + assert minter.calls == [("u1", "team-b")] + fresh = _Redeemer(ASSERTED.model_copy(update={"replay_key": "mcp_gateway_dcr:id_jag:jti-2", "scope": None})) + again = await _redeem_assertion_native(minter, fresh, cache=cache) + assert again.status_code == 200 + assert "scope" not in json.loads(again.body) + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_claims_the_jti_before_it_mints(): + """A mint the gateway could not complete still consumed the assertion: the client gets the + 503 and comes back with a fresh assertion, never with the same one.""" + minter, cache = _Minter(result="unavailable"), DualCache() + first = await _redeem_assertion_native(minter, _Redeemer(), cache=cache) + assert first.status_code == 503 + retried = await _redeem_assertion_native(_Minter(), _Redeemer(), cache=cache) + assert retried.status_code == 400 + assert json.loads(retried.body)["error"] == "invalid_grant" + + +@pytest.mark.parametrize( + "error, status", + [ + ("invalid_client", 401), + ("invalid_grant", 400), + ("invalid_request", 400), + ("unsupported_grant_type", 400), + ("temporarily_unavailable", 503), + ], +) +@pytest.mark.asyncio +async def test_jwt_bearer_grant_refusals_keep_their_oauth_error_and_status(error, status): + minter, cache = _Minter(), DualCache() + refusal = SubjectTokenRefusal(error=error, description="why") + response = await _redeem_assertion_native(minter, _Redeemer(refusal), cache=cache) + assert response.status_code == status + assert json.loads(response.body) == {"error": error, "error_description": "why"} + assert minter.calls == [] + served = await _redeem_assertion_native(minter, _Redeemer(), cache=cache) + assert served.status_code == 200 + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_requires_an_assertion_and_the_proxy_api_target(): + minter, redeemer = _Minter(), _Redeemer() + missing = await _redeem_assertion_native(minter, redeemer, assertion=None) + assert missing.status_code == 400 + assert json.loads(missing.body)["error"] == "invalid_request" + foreign = await _redeem_assertion_native(minter, redeemer, resource="https://elsewhere.example/api") + assert foreign.status_code == 400 + assert json.loads(foreign.body)["error"] == "invalid_target" + assert redeemer.calls == [] + assert minter.calls == [] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_is_refused_where_no_redeemer_is_wired(): + minter = _Minter() + response = await _redeem_native( + None, REQUESTING_APP, minter, grant_type=JWT_BEARER_GRANT_TYPE, assertion=ID_JAG, client_secret=None + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert minter.calls == [] + + +@pytest.mark.asyncio +async def test_an_unknown_grant_type_names_the_jwt_bearer_grant_among_the_served_ones(): + response = await _redeem_native(None, REQUESTING_APP, _Minter(), grant_type="urn:example:nope") + assert response.status_code == 400 + assert JWT_BEARER_GRANT_TYPE in json.loads(response.body)["error_description"] + + +@pytest.mark.parametrize( + "token_exchange, jwt_bearer, expected", + [ + (False, False, ("authorization_code", "refresh_token")), + (True, False, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)), + (False, True, ("authorization_code", "refresh_token", JWT_BEARER_GRANT_TYPE)), + (True, True, ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE, JWT_BEARER_GRANT_TYPE)), + ], +) +def test_the_jwt_bearer_grant_is_advertised_exactly_when_it_is_served(token_exchange, jwt_bearer, expected): + assert supported_grant_types(token_exchange, jwt_bearer) == expected + contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), token_exchange, jwt_bearer) + assert contract["grant_types_supported"] == expected + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_narrows_the_credential_to_the_routes_the_scope_bought(): + """The route groups the assertion's scope resolved to reach the minter, so the credential is + enforced on every proxy route the way a virtual key's allowed_routes is, and the response states + exactly the scope that was enforced. No scope on the assertion mints the full credential.""" + minter = _Minter() + both = ASSERTED.model_copy( + update={"scope": "litellm:mcp litellm:llm", "allowed_routes": ("mcp_inference_routes", "openai_routes")} + ) + scoped = await _redeem_assertion_native(minter, _Redeemer(both), cache=DualCache()) + assert scoped.status_code == 200 + assert json.loads(scoped.body)["scope"] == "litellm:mcp litellm:llm" + full = await _redeem_assertion_native( + minter, _Redeemer(ASSERTED.model_copy(update={"scope": None, "allowed_routes": None})), cache=DualCache() + ) + assert full.status_code == 200 + assert "scope" not in json.loads(full.body) + assert minter.allowed_routes == [("mcp_inference_routes", "openai_routes"), None] + + +@pytest.mark.asyncio +async def test_jwt_bearer_grant_accepts_the_mcp_resource_indicator_as_its_target(): + """RFC 8707 clients send the ID-JAG's audience as ``resource``, and that audience may be the MCP + resource ``{base}/mcp``, so it is a valid target next to the proxy origin. Anything else stays + ``invalid_target``.""" + minter = _Minter() + served = await _redeem_assertion_native( + minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/mcp" + ) + assert served.status_code == 200 + foreign = await _redeem_assertion_native( + minter, _Redeemer(), cache=DualCache(), resource=PROXY_API_RESOURCE.rstrip("/") + "/other/mcp" + ) + assert json.loads(foreign.body)["error"] == "invalid_target" + assert minter.calls == [("u1", "team-b")] + + +@pytest.mark.parametrize("jwt_bearer_available", [True, False]) +def test_aggregate_metadata_advertises_basic_client_auth_only_with_the_grant_that_reads_it(jwt_bearer_available): + """Only the jwt-bearer redemption reads an HTTP Basic client secret, the other grants take + ``client_secret`` from the form, so ``client_secret_basic`` is advertised only when that grant + is on. A client following the document otherwise sends Basic alone and authenticates with no + secret.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_authorization_server_response, + ) + + document = _build_aggregate_authorization_server_response(_request(), True, jwt_bearer_available) + methods = document["token_endpoint_auth_methods_supported"] + assert isinstance(methods, list) + assert ("client_secret_basic" in methods) is jwt_bearer_available + assert "none" in methods and "client_secret_post" in methods diff --git a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py new file mode 100644 index 00000000000..fdfa0544ad1 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py @@ -0,0 +1,738 @@ +import hashlib +import logging +import math +from base64 import b64encode +from collections.abc import Mapping, Sequence +from typing import Final + +import jwt +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + AssertionIdentity, + ConsentTeam, + ReloadUserFailure, + SubjectIdentity, + SubjectTokenRefusal, +) +from litellm.proxy._experimental.mcp_server.id_jag_grant import ( + ID_JAG_REQUIRED_CLAIMS, + REJECTED_ASSERTION, + TEAM_HEADER, + AssertionRedemption, + GrantedScope, + IdJagPrerequisites, + LookupTeams, + _claim_reader, + basic_client_id, + gateway_audiences, + granted_scope, + identity_from_id_jag, + jwt_bearer_available, + names_mcp_resource, + proxy_user_lookup, + select_team, + subject_resolver, +) +from litellm.proxy._types import IdJagRequestingApp, LiteLLM_JWTAuth, LiteLLMRoutes +from litellm.proxy.auth.handle_jwt import ( + ID_JAG_TYP, + AssertionRejection, + JWKSUnreachableError, + JWTHandler, + jwks_unavailable_exception, +) + +ISSUER = "https://idp.example.com" +GATEWAY = "https://gateway.example.com" +AUDIENCES = (f"{GATEWAY}/mcp", GATEWAY) +NOW = 1_800_000_000.0 +PUBLIC_APP = IdJagRequestingApp(client_id="0oa-public") +CONFIDENTIAL_APP = IdJagRequestingApp(client_id="0oa-confidential", client_secret="s3cret") +SPACED_APP = IdJagRequestingApp(client_id="0oa spaced", client_secret="app secret") +EVERY_GATE_HOLDS = { + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "maps_jwts_to_virtual_keys": False, + "shared_replay_store": True, + "base_url": "https://gateway.example.com", + "requesting_apps": (PUBLIC_APP, CONFIDENTIAL_APP, SPACED_APP), +} +JWKS_DOWN = jwks_unavailable_exception(JWKSUnreachableError("ConnectError fetching the IdP's JWKS after 3 attempts")) + + +def _claims(**overrides: object) -> dict[str, object]: + return { + "iss": ISSUER, + "sub": "okta-user-1", + "aud": AUDIENCES[1], + "exp": NOW + 300, + "iat": NOW - 5, + "jti": "jti-1", + "client_id": PUBLIC_APP.client_id, + "scope": "litellm:mcp", + **overrides, + } + + +def _assertion(typ: str = ID_JAG_TYP, **claims: object) -> str: + return jwt.encode( + _claims(**claims), + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": typ}, + ) + + +class _Verifier: + """Stands in for the injected ``DecodeAssertion`` port; records each call as a new tuple.""" + + def __init__( + self, result: Mapping[str, object] | AssertionRejection | None = None, raises: BaseException | None = None + ) -> None: + self.calls: tuple[tuple[str, tuple[str, ...], tuple[str, ...]], ...] = () + self.result = result + self.raises = raises + + async def __call__( + self, assertion: str, audience: Sequence[str], required_claims: Sequence[str] + ) -> Mapping[str, object] | AssertionRejection: + self.calls = (*self.calls, (assertion, tuple(audience), tuple(required_claims))) # rebind-ok: immutable record + if self.raises is not None: + raise self.raises + return self.result if self.result is not None else jwt.decode(assertion, options={"verify_signature": False}) + + +class _Resolver: + """Stands in for the injected ``ResolveSubject`` port; records each call as a new tuple.""" + + def __init__( + self, + user_id: str | None = "u1", + raises: BaseException | None = None, + result: SubjectIdentity | SubjectTokenRefusal | None = None, + ) -> None: + self.calls: tuple[tuple[Mapping[str, object], str | None], ...] = () + self.user_id = user_id + self.raises = raises + self.result = result + + async def __call__( + self, claims: Mapping[str, object], team_header: str | None + ) -> SubjectIdentity | SubjectTokenRefusal | None: + self.calls = (*self.calls, (dict(claims), team_header)) # rebind-ok: immutable record + if self.raises is not None: + raise self.raises + if self.result is not None: + return self.result + return None if self.user_id is None else SubjectIdentity(user_id=self.user_id, team_id=team_header) + + +def _redemption( + assertion: str, + client_id: str = PUBLIC_APP.client_id, + client_secret: str | None = None, + authorization: str | None = None, + team_id: str | None = "team-b", + leeway: float = 0.0, +) -> AssertionRedemption: + return AssertionRedemption( + assertion=assertion, + client_id=client_id, + client_secret=client_secret, + authorization_header=authorization, + team_id=team_id, + audiences=AUDIENCES, + now=NOW, + leeway_seconds=leeway, + ) + + +@pytest.mark.asyncio +async def test_the_replay_marker_outlives_the_leeway_the_verifier_grants_exp(): + """PyJWT accepts an assertion until ``exp`` plus the handler's leeway, so the marker must + live that long too, including for an assertion already past ``exp`` but inside the leeway.""" + within_leeway = await _identity(_redemption(_assertion(), leeway=30), _Verifier(), _Resolver()) + assert isinstance(within_leeway, AssertionIdentity) + assert within_leeway.replay_ttl_seconds == 330 + just_expired = await _identity( + _redemption(_assertion(), leeway=30), _Verifier(result=_claims(exp=NOW - 10)), _Resolver() + ) + assert isinstance(just_expired, AssertionIdentity) + assert just_expired.replay_ttl_seconds == 20 + + +async def _identity( + redemption: AssertionRedemption, + verifier: _Verifier | None = None, + resolver: _Resolver | None = None, + **unmet: object, +) -> AssertionIdentity | SubjectTokenRefusal: + return await identity_from_id_jag( + redemption, + prerequisites=IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}), + decode=verifier if verifier is not None else _Verifier(), + resolve_subject=resolver if resolver is not None else _Resolver(), + ) + + +@pytest.mark.asyncio +async def test_a_valid_id_jag_names_the_user_its_single_use_key_and_its_scope(): + """The verifier is asked for the gateway's own issuer values as audience and for every claim + the ID-JAG profile requires, the subject lookup sees the verified claims, and the identity + carries the team asked for, a replay key derived from iss and jti, a marker TTL that outlives + the assertion, and the assertion's scope.""" + assertion, verifier, resolver = _assertion(), _Verifier(), _Resolver() + identity = await _identity(_redemption(assertion), verifier, resolver) + assert isinstance(identity, AssertionIdentity) + assert verifier.calls == ((assertion, AUDIENCES, ID_JAG_REQUIRED_CLAIMS),) + assert {"iss", "sub", "aud", "exp", "iat", "jti", "client_id"} <= set(ID_JAG_REQUIRED_CLAIMS) + assert resolver.calls == ((_claims(), "team-b"),) + assert (identity.user_id, identity.team_id, identity.scope) == ("u1", "team-b", "litellm:mcp") + assert identity.allowed_routes == ("mcp_inference_routes",) + assert identity.replay_key == "mcp_gateway_dcr:id_jag:" + hashlib.sha256(f"{ISSUER}\njti-1".encode()).hexdigest() + assert identity.replay_ttl_seconds == math.ceil(300) + other_jti = await _identity(_redemption(_assertion(jti="jti-2")), _Verifier(), _Resolver()) + assert isinstance(other_jti, AssertionIdentity) + assert other_jti.replay_key != identity.replay_key + teamless = await _identity(_redemption(_assertion(scope=""), team_id=None)) + assert isinstance(teamless, AssertionIdentity) + assert (teamless.team_id, teamless.scope, teamless.allowed_routes) == (None, None, None) + + +@pytest.mark.asyncio +async def test_a_requesting_app_the_operator_did_not_allowlist_learns_nothing_about_the_assertion(): + verifier, resolver = _Verifier(), _Resolver() + refused = await _identity(_redemption(_assertion(), client_id="0oa-stranger"), verifier, resolver) + assert refused == SubjectTokenRefusal( + error="invalid_client", description="client_id is not an allowlisted requesting app" + ) + assert verifier.calls == () + assert resolver.calls == () + + +def _basic(client_id: str, secret: str) -> str: + return "Basic " + b64encode(f"{client_id}:{secret}".encode()).decode() + + +@pytest.mark.parametrize( + "client_secret, authorization, accepted", + [ + (None, None, False), + ("wrong", None, False), + ("s3cr\u00e9t", None, False), + ("s3cret", None, True), + (None, _basic(CONFIDENTIAL_APP.client_id, "s3cr\u00e9t"), False), + (None, _basic(CONFIDENTIAL_APP.client_id, "s3cret"), True), + (None, _basic(CONFIDENTIAL_APP.client_id, "wrong"), False), + (None, _basic(PUBLIC_APP.client_id, "s3cret"), False), + ("s3cret", _basic(CONFIDENTIAL_APP.client_id, "wrong"), False), + (None, "Basic not-base64!", False), + (None, "Bearer s3cret", False), + ], +) +@pytest.mark.asyncio +async def test_a_confidential_app_must_authenticate_with_its_secret(client_secret, authorization, accepted): + verifier = _Verifier() + assertion = _assertion(client_id=CONFIDENTIAL_APP.client_id) + redemption = _redemption(assertion, CONFIDENTIAL_APP.client_id, client_secret, authorization) + identity = await _identity(redemption, verifier) + if accepted: + assert isinstance(identity, AssertionIdentity) + return + assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + assert verifier.calls == () + + +@pytest.mark.parametrize( + "client_secret, authorization, accepted", + [ + ("app secret", None, True), + (None, "Basic " + b64encode(b"0oa+spaced:app+secret").decode(), True), + (None, "Basic " + b64encode(b"0oa%20spaced:app%20secret").decode(), True), + (None, "Basic " + b64encode(b"0oa spaced:app secret").decode(), True), + (None, "Basic " + b64encode(b"0oa+spaced:app+wrong").decode(), False), + (None, "Basic " + b64encode(b"0oa-spaced:app+secret").decode(), False), + ], +) +@pytest.mark.asyncio +async def test_basic_credentials_are_form_decoded_so_a_space_in_an_id_or_secret_still_matches( + client_secret, authorization, accepted +): + """RFC 6749 section 2.3.1 form-encodes the id and secret before base64, so ``+`` is a space. The + endpoint hands over the ``client_id`` it read from the header, so the header must name the same app.""" + assertion = _assertion(client_id=SPACED_APP.client_id) + identity = await _identity(_redemption(assertion, SPACED_APP.client_id, client_secret, authorization), _Verifier()) + if accepted: + assert isinstance(identity, AssertionIdentity) + return + assert identity == SubjectTokenRefusal(error="invalid_client", description="client authentication failed") + + +@pytest.mark.asyncio +async def test_a_public_app_is_identified_by_client_id_alone(): + identity = await _identity(_redemption(_assertion(), client_secret="anything", authorization="Bearer x")) + assert isinstance(identity, AssertionIdentity) + + +@pytest.mark.parametrize( + "assertion, description", + [ + ("not-a-jwt", "assertion is not a JWT"), + (_assertion(typ="JWT"), f"assertion typ must be {ID_JAG_TYP}"), + (_assertion(typ="at+jwt"), f"assertion typ must be {ID_JAG_TYP}"), + ], +) +@pytest.mark.asyncio +async def test_a_bearer_or_id_token_is_not_an_id_jag(assertion, description): + verifier = _Verifier() + refused = await _identity(_redemption(assertion), verifier) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=description) + assert verifier.calls == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verifier, reason", + [ + (_Verifier(result=AssertionRejection(reason="Signature verification failed")), "Signature verification"), + (_Verifier(result=AssertionRejection(reason="Audience doesn't match")), "Audience"), + (_Verifier(result=AssertionRejection(reason='Token is missing the "jti" claim')), "jti"), + (_Verifier(result=AssertionRejection(reason="Signature has expired")), "expired"), + (_Verifier(result=AssertionRejection(reason="iss is not an issuer this gateway trusts")), "iss"), + (_Verifier(raises=Exception("Validation fails: Invalid audience")), "Invalid audience"), + (_Verifier(result=AssertionRejection(reason="bad kid\nWARNING forged line")), "bad kid WARNING forged line"), + ], +) +async def test_an_assertion_the_verifier_rejects_is_refused_with_the_reason_kept_in_the_log(verifier, reason, caplog): + resolver = _Resolver() + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), verifier, resolver) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert reason not in refused.description + assert reason in caplog.text + assert resolver.calls == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "claims, reason", + [ + ({"client_id": "0oa-confidential"}, "client_id"), + ({"exp": NOW}, "expired"), + ({"exp": NOW - 1}, "expired"), + ({"jti": ""}, "jti"), + ({"jti": 7}, "jti"), + ({"iss": None}, "iss"), + ({"cnf": {"jkt": "thumbprint"}}, "proof"), + ], +) +async def test_verified_claims_that_do_not_fit_the_profile_are_refused(claims, reason, caplog): + resolver = _Resolver() + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(**claims)), resolver) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert reason in caplog.text + assert resolver.calls == () + + +@pytest.mark.asyncio +async def test_an_assertion_that_names_no_user_the_gateway_knows_cannot_be_redeemed(caplog): + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + refused = await _identity(_redemption(_assertion()), resolver=_Resolver(user_id=None)) + assert refused == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + assert "names no user" in caplog.text + missing = await _identity( + _redemption(_assertion()), resolver=_Resolver(raises=ValueError("User doesn't exist in db. 'user_id'=u1")) + ) + assert missing == SubjectTokenRefusal(error="invalid_grant", description=REJECTED_ASSERTION) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verifier, resolver", + [ + (_Verifier(raises=JWKS_DOWN), _Resolver()), + (_Verifier(), _Resolver(raises=HTTPException(status_code=503, detail="database unavailable"))), + ], +) +async def test_a_gateway_outage_is_reported_as_retryable_not_as_a_bad_assertion(verifier, resolver): + refused = await _identity(_redemption(_assertion()), verifier, resolver) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "temporarily_unavailable" + assert "retry" in refused.description + assert "assertion" in refused.description + assert "subject_token" not in refused.description + + +@pytest.mark.parametrize( + "unmet, description", + [ + ({"jwt_auth_enabled": False}, "JWT auth is not enabled"), + ({"has_database": False}, "no database"), + ({"licensed": False}, "enterprise only"), + ({"maps_jwts_to_virtual_keys": True}, "maps IdP tokens to virtual keys"), + ({"shared_replay_store": False}, "Redis"), + ({"base_url": None}, "PROXY_BASE_URL"), + ({"requesting_apps": ()}, "no requesting apps allowlisted"), + ], +) +@pytest.mark.asyncio +async def test_the_deployment_gates_refuse_before_any_verification(unmet, description): + verifier = _Verifier() + refused = await _identity(_redemption(_assertion()), verifier, **unmet) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "unsupported_grant_type" + assert description in refused.description + assert verifier.calls == () + assert IdJagPrerequisites(**{**EVERY_GATE_HOLDS, **unmet}).available is False + assert IdJagPrerequisites(**EVERY_GATE_HOLDS).available is True + + +def _running_jwt_handler(litellm_jwtauth: LiteLLM_JWTAuth | None) -> JWTHandler: + handler: Final = JWTHandler() + if litellm_jwtauth is not None: + handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth) + return handler + + +@pytest.mark.parametrize( + "general_settings, prisma_client, premium_user, litellm_jwtauth, expected", + [ + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), True), + ( + {"enable_jwt_auth": True, "redis": None}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), + False, + ), + ( + {"enable_jwt_auth": True, "base_url": ""}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), + False, + ), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[]), False), + ( + {"enable_jwt_auth": True}, + object(), + True, + LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP], virtual_key_claim_field="vkey"), + False, + ), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), True, None, False), + ({}, object(), True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(id_jag_requesting_apps=[PUBLIC_APP]), False), + ], +) +def test_availability_is_read_from_the_running_proxy( + monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected +): + monkeypatch.setenv("PROXY_BASE_URL", general_settings.get("base_url", "https://gateway.example.com")) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", general_settings.get("redis", object())) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth)) + assert jwt_bearer_available() is expected + + +def test_the_assertion_audience_is_the_fixed_gateway_url_not_the_request_host(): + """An ID-JAG must name this gateway's MCP resource or its authorization server under the + operator-fixed PROXY_BASE_URL. Without it there is no audience, so there is no grant.""" + assert gateway_audiences("https://gateway.example.com") == ( + "https://gateway.example.com/mcp", + "https://gateway.example.com", + ) + assert gateway_audiences(None) == () + + +@pytest.mark.asyncio +async def test_a_team_the_subject_lookup_cannot_settle_is_refused_before_the_assertion_is_consumed(): + """The resolver's refusal (a user on several teams who named none, or a header naming a team + they are not on) comes back as is, so the client fixes its request and presents the same + assertion again instead of finding its jti burned.""" + refusal = SubjectTokenRefusal(error="invalid_request", description=f"send the {TEAM_HEADER} header") + refused = await _identity(_redemption(_assertion()), resolver=_Resolver(result=refusal)) + assert refused is refusal + teamed = await _identity( + _redemption(_assertion()), resolver=_Resolver(result=SubjectIdentity(user_id="u1", team_id="t9")) + ) + assert isinstance(teamed, AssertionIdentity) + assert teamed.team_id == "t9" + + +ENG = ConsentTeam(team_id="team-eng", team_alias="engineering") +OPS = ConsentTeam(team_id="team-ops", team_alias=None) + + +@pytest.mark.parametrize( + "team_header, claim_teams, teams, expected", + [ + (None, (), (), None), + (None, (), (ENG,), "team-eng"), + ("team-eng", (), (ENG, OPS), "team-eng"), + ("engineering", (), (ENG, OPS), "team-eng"), + ("team-ops", ("team-eng",), (ENG, OPS), "team-ops"), + (None, ("team-ops",), (ENG, OPS), "team-ops"), + (None, ("team-sales", "team-ops"), (ENG, OPS), "team-ops"), + (None, ("0oa-xaa-public-app",), (ENG,), "team-eng"), + (None, ("0oa-xaa-public-app",), (), None), + ], +) +def test_select_team_prefers_the_header_then_the_assertion_then_the_only_team( + team_header, claim_teams, teams, expected +): + assert select_team(team_header, claim_teams, teams) == expected + + +class _Reader: + def __init__(self, user_id=None, email=None, team_ids=(), allowed_domain=None): + self.user_id, self.email, self.team_ids, self.allowed_domain = user_id, email, team_ids, allowed_domain + + def is_allowed_domain(self, user_email): + return self.allowed_domain is None or user_email.split("@")[-1] == self.allowed_domain + + def get_user_id(self, token, default_value): + return self.user_id or default_value + + def get_user_email(self, token, default_value): + return self.email or default_value + + def get_all_jwt_team_ids(self, token): + return list(self.team_ids) + + +async def _known_user(user_id: str, user_email: str | None) -> str | None: + return user_id if user_id in {"okta-user-1", "u-mapped"} else None + + +async def _two_teams(user_id: str) -> tuple[ConsentTeam, ...]: + return (ENG, OPS) + + +async def _teams_unreadable(user_id: str) -> ReloadUserFailure: + return "unavailable" + + +def _teams_failing(failure: ReloadUserFailure) -> LookupTeams: + async def lookup_teams(user_id: str) -> tuple[ConsentTeam, ...] | ReloadUserFailure: + return failure + + return lookup_teams + + +@pytest.mark.asyncio +async def test_the_subject_resolver_binds_the_claims_to_a_known_user_and_their_asserted_team(): + """The user comes from the JWT auth user mapping, else ``sub``, and must exist. The team the + assertion asserts settles a multi-team user, the header still wins, and unreadable teams are a + retryable refusal rather than a credential for the wrong team.""" + asserted = subject_resolver(_Reader(team_ids=("team-ops",)), _known_user, _two_teams) + assert await asserted(_claims(), None) == SubjectIdentity(user_id="okta-user-1", team_id="team-ops") + assert await asserted(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng") + mapped = subject_resolver(_Reader(user_id="u-mapped", email="m@example.com"), _known_user, _two_teams) + assert await mapped(_claims(), "engineering") == SubjectIdentity(user_id="u-mapped", team_id="team-eng") + unknown = subject_resolver(_Reader(), _known_user, _two_teams) + assert await unknown(_claims(sub="nobody"), None) is None + unreadable = await subject_resolver(_Reader(), _known_user, _teams_unreadable)(_claims(), None) + assert isinstance(unreadable, SubjectTokenRefusal) + assert unreadable.error == "temporarily_unavailable" + + +@pytest.mark.parametrize("failure", ["no_active_key", "unresolvable"]) +@pytest.mark.asyncio +async def test_a_blocked_or_unmapped_user_is_a_bad_grant_not_an_outage_to_retry(failure): + """A deactivated user, or one the resolver cannot map, can retry forever and never succeed, so + the refusal is invalid_grant. Only a database the gateway could not reach is retryable.""" + refused = await subject_resolver(_Reader(), _known_user, _teams_failing(failure))(_claims(), None) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_grant" + assert "retry" not in refused.description + outage = await subject_resolver(_Reader(), _known_user, _teams_failing("faulted"))(_claims(), None) + assert isinstance(outage, SubjectTokenRefusal) + assert outage.error == "temporarily_unavailable" + + +def test_the_assertion_team_claims_ignore_team_id_default(): + """``team_id_default`` is how bearer JWT auth picks a team for a request, not a claim the + assertion made, so a user outside that team must still get their only team or no team.""" + handler = JWTHandler() + handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_default="default-team") + claim_teams = tuple(_claim_reader(handler).get_all_jwt_team_ids(token={"iss": ISSUER, "sub": "u1"})) + assert claim_teams == () + assert select_team(None, claim_teams, (ENG,)) == "team-eng" + assert select_team(None, claim_teams, ()) is None + + +@pytest.mark.parametrize( + "team_header, claim_teams, teams, error, mention", + [ + ("team-sales", (), (ENG, OPS), "invalid_request", TEAM_HEADER), + ("engineering", (), (), "invalid_request", TEAM_HEADER), + (None, (), (ENG, OPS), "invalid_request", TEAM_HEADER), + (None, ("team-eng", "team-ops"), (ENG, OPS), "invalid_request", TEAM_HEADER), + (None, ("0oa-xaa-public-app",), (ENG, OPS), "invalid_request", TEAM_HEADER), + ], +) +def test_select_team_refuses_a_team_the_user_is_not_on_or_an_unnamed_choice( + team_header, claim_teams, teams, error, mention +): + refused = select_team(team_header, claim_teams, teams) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == error + assert mention in refused.description + + +@pytest.mark.parametrize( + "authorization, expected", + [ + (None, None), + ("Bearer abc", None), + ("Basic not-base64!", None), + (_basic("0oa-confidential", "s3cret"), "0oa-confidential"), + ("basic " + b64encode(b"0oa%20x:s3cret").decode(), "0oa x"), + ("Basic " + b64encode(b"0oa+x:s3cret").decode(), "0oa x"), + ("Basic " + b64encode(b"no-separator").decode(), None), + ], +) +def test_basic_client_id_reads_the_app_from_the_authorization_header(authorization, expected): + assert basic_client_id(authorization) == expected + + +MCP_ROUTE_GROUPS = ("mcp_inference_routes",) +LLM_ROUTE_GROUPS = ("llm_api_routes_without_mcp",) + + +@pytest.mark.parametrize( + "claim, scope, allowed_routes", + [ + (None, None, None), + ("", None, None), + ("litellm:mcp", "litellm:mcp", MCP_ROUTE_GROUPS), + ("litellm:llm", "litellm:llm", LLM_ROUTE_GROUPS), + ("litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS), + ("litellm:llm litellm:mcp litellm:llm", "litellm:mcp litellm:llm", MCP_ROUTE_GROUPS + LLM_ROUTE_GROUPS), + ], +) +def test_the_scope_claim_buys_exactly_the_route_groups_it_names(claim, scope, allowed_routes): + """``litellm:mcp`` is the MCP gateway's data-plane routes, ``litellm:llm`` the LLM API routes + without their MCP members, both together are both, and no claim is the full credential. The + scope echoed back is canonical: deduplicated, in vocabulary order, so it is exactly what is + enforced. Every group named must be one the route check resolves.""" + granted = granted_scope(claim, mcp_resource=False) + assert (granted.scope, granted.allowed_routes) == (scope, allowed_routes) + if allowed_routes is not None: + assert set(allowed_routes) <= set(LiteLLMRoutes._member_names_) + mcp_members = set(LiteLLMRoutes.mcp_inference_routes.value) + llm_members = set().union(*(LiteLLMRoutes[group].value for group in LLM_ROUTE_GROUPS)) + assert "/chat/completions" in llm_members and "/mcp" not in llm_members + assert "/mcp" in mcp_members and "/chat/completions" not in mcp_members + + +@pytest.mark.parametrize("claim", ["openid", "litellm:mcp openid", "litellm:admin", 7, ["litellm:mcp"]]) +@pytest.mark.asyncio +async def test_a_scope_the_gateway_does_not_know_is_refused_as_invalid_scope_before_any_lookup(claim): + refusal = granted_scope(claim, mcp_resource=False) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_scope" + assert "litellm:mcp" in refusal.description and "litellm:llm" in refusal.description + resolver = _Resolver() + refused = await _identity(_redemption(_assertion()), _Verifier(result=_claims(scope=claim)), resolver) + assert refused == refusal + assert resolver.calls == () + + +@pytest.mark.asyncio +async def test_a_linked_sso_account_resolves_by_its_provider_subject_without_an_email_claim(): + """An SSO login stores the provider subject as ``sso_user_id`` on a user whose ``user_id`` is + something else, so the lookup must offer the subject as both, with no email to fall back on.""" + from litellm.models.user import LiteLLM_UserTable + + async def get_user(*, user_id, user_email, sso_user_id): + if sso_user_id == "00u-okta-subject" and user_email is None: + return LiteLLM_UserTable(user_id="local-user-7", user_email=None) + return None + + lookup = proxy_user_lookup(get_user) + assert await lookup("00u-okta-subject", None) == "local-user-7" + assert await lookup("nobody", None) is None + + +@pytest.mark.parametrize("email", ["eve@evil.example", None]) +@pytest.mark.asyncio +async def test_the_allowed_email_domain_gates_the_assertions_subject_like_bearer_jwt_auth(email): + """With ``user_allowed_email_domain`` set, bearer JWT auth refuses an identity whose email is + outside the domain or missing. An assertion for an existing user gets the same refusal before + any lookup, as invalid_grant, and a matching email still resolves.""" + refused = await subject_resolver(_Reader(email=email, allowed_domain="corp.example"), _known_user, _two_teams)( + _claims(), "team-eng" + ) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_grant" + allowed = subject_resolver( + _Reader(email="ann@corp.example", allowed_domain="corp.example"), _known_user, _two_teams + ) + assert await allowed(_claims(), "team-eng") == SubjectIdentity(user_id="okta-user-1", team_id="team-eng") + + +@pytest.mark.parametrize("claim", [None, "", "litellm:mcp", "litellm:mcp litellm:mcp"]) +def test_an_assertion_for_the_mcp_resource_is_ceilinged_to_the_mcp_routes(claim): + """``aud = {base}/mcp`` means the identity provider named the MCP resource, so no scope or + ``litellm:mcp`` both mint an MCP-only credential and echo ``litellm:mcp``.""" + granted = granted_scope(claim, mcp_resource=True) + assert granted == GrantedScope(scope="litellm:mcp", allowed_routes=("mcp_inference_routes",)) + + +@pytest.mark.parametrize("claim", ["litellm:llm", "litellm:mcp litellm:llm"]) +def test_llm_routes_cannot_be_granted_under_the_mcp_resource_audience(claim): + refusal = granted_scope(claim, mcp_resource=True) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_scope" + assert "/mcp audience" in refusal.description and "bare origin" in refusal.description + + +@pytest.mark.parametrize( + "aud, expected", + [ + (AUDIENCES[0], True), + (AUDIENCES[1], False), + ([AUDIENCES[0]], True), + ([AUDIENCES[0], AUDIENCES[1]], False), + ("https://other.example", False), + (None, False), + ], +) +def test_the_mcp_resource_audience_is_the_mcp_path_without_the_bare_origin(aud, expected): + assert names_mcp_resource(aud, GATEWAY) is expected + assert names_mcp_resource(aud, None) is False + + +@pytest.mark.asyncio +async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_llm_scope(): + """End to end through identity_from_id_jag: the same verified claims, with ``aud`` the MCP + resource, come out ceilinged, and asking for LLM routes under that audience is a scope error.""" + narrowed = await _identity(_redemption(_assertion(aud=AUDIENCES[0])), _Verifier(), _Resolver()) + assert isinstance(narrowed, AssertionIdentity) + assert (narrowed.scope, narrowed.allowed_routes) == ("litellm:mcp", ("mcp_inference_routes",)) + refused = await _identity(_redemption(_assertion(aud=AUDIENCES[0], scope="litellm:llm")), _Verifier(), _Resolver()) + assert isinstance(refused, SubjectTokenRefusal) + assert refused.error == "invalid_scope" + unscoped = jwt.encode( + {key: value for key, value in _claims(aud=AUDIENCES[1]).items() if key != "scope"}, + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": ID_JAG_TYP}, + ) + full = await _identity(_redemption(unscoped), _Verifier(), _Resolver()) + assert isinstance(full, AssertionIdentity) + assert (full.scope, full.allowed_routes) == (None, None) + scoped = await _identity(_redemption(_assertion(aud=AUDIENCES[1], scope="litellm:llm")), _Verifier(), _Resolver()) + assert isinstance(scoped, AssertionIdentity) + assert scoped.scope == "litellm:llm" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 536c1eb0b44..40901d1dd96 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1641,9 +1641,10 @@ class TestMCPServerManager: assert server.oauth2_flow is None @pytest.mark.asyncio - async def test_load_servers_from_config_warns_for_id_jag_with_google_sso(self, monkeypatch, caplog): + async def test_load_servers_from_config_is_silent_for_id_jag_with_saml_sso(self, monkeypatch, caplog): + """A SAML login retains the signed assertion as the saml2 subject token, so SAML is not a gap.""" self._clear_sso_env(monkeypatch) - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") manager = MCPServerManager() with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), @@ -1652,10 +1653,25 @@ class TestMCPServerManager: ): await manager.load_servers_from_config(self._id_jag_config()) - warnings = [message for message in caplog.messages if "oauth2_id_jag" in message] - assert len(warnings) == 1 - assert "idjag_server" in warnings[0] - assert "GENERIC_CLIENT_ID" in warnings[0] + assert [message for message in caplog.messages if "oauth2_id_jag" in message] == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + async def test_load_servers_from_config_is_silent_for_id_jag_under_oidc_logins( + self, monkeypatch, caplog, client_id_env + ): + """Every OIDC login captures the IdP id_token an id_jag server spends, so none of them is a gap.""" + self._clear_sso_env(monkeypatch) + monkeypatch.setenv(client_id_env, "cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) @pytest.mark.asyncio async def test_load_servers_from_config_does_not_warn_for_id_jag_without_sso(self, monkeypatch, caplog): diff --git a/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 8484dfdde72..f39aec82c3e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -1,15 +1,21 @@ """Tests for minting the ``lite login`` credential from a consented native-client grant.""" -from unittest.mock import ANY, AsyncMock, MagicMock +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential -from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential +from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( + lookup_consent_teams, + lookup_consent_teams_from_database, + mint_proxy_credential, +) from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail @@ -235,3 +241,82 @@ async def test_lookup_consent_teams_passes_user_lookup_failures_through(failure, async def test_lookup_consent_teams_reports_unavailable_when_details_cannot_load(load_user, fetch_teams): fetch_teams.return_value = None assert await lookup_consent_teams("u1") == "unavailable" + + +def _route_status(minted: MintedProxyCredential, route: str, method: str | None = None) -> int: + request = None if method is None else MagicMock(method=method) + try: + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=_decoded(minted), request=request) + except HTTPException as refused: + return refused.status_code + return 200 + + +@pytest.mark.parametrize( + "allowed_routes, chat_completions, mcp", + [ + (None, True, True), + (("mcp_inference_routes",), False, True), + (("llm_api_routes_without_mcp",), True, False), + (("mcp_inference_routes", "llm_api_routes_without_mcp"), True, True), + ], + ids=["full credential", "mcp only", "llm only", "both"], +) +@pytest.mark.asyncio +async def test_mint_narrows_the_credential_to_the_allowed_routes_the_route_check_enforces( + load_user, fetch_teams, allowed_routes, chat_completions, mcp +): + """The allowed_routes a grant's scope resolved to ride inside the credential, so the same route + check a virtual key goes through on every request refuses /chat/completions for an MCP-only + credential and /mcp for an LLM-only one, and a credential minted without them is unchanged.""" + minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes) + assert isinstance(minted, MintedProxyCredential) + assert _decoded(minted).allowed_routes == list(allowed_routes or ()) + chat_status, mcp_status = (200 if chat_completions else 403), (200 if mcp else 403) + assert _route_status(minted, "/chat/completions") == chat_status + assert _route_status(minted, "/v1/chat/completions") == chat_status + assert _route_status(minted, "/mcp") == mcp_status + assert _route_status(minted, "/mcp/tools/call") == mcp_status + + +@pytest.mark.asyncio +async def test_lookup_consent_teams_from_database_reads_the_row_the_minter_will_recheck(load_user, fetch_teams): + """A grant consumed once a team is picked must pick it from the same membership the minter re-reads, + or a stale cache would burn the assertion on a request the caller could have fixed.""" + teams = await lookup_consent_teams_from_database("u1") + load_user.assert_awaited_once_with("u1", source="database") + assert teams == (ConsentTeam(team_id="team-a", team_alias="Team A"), ConsentTeam(team_id="team-b")) + + +@pytest.mark.parametrize( + "allowed_routes, registered_pass_through, auto_router_session, mcp_server_listing, agent_registry", + [ + (("llm_api_routes",), 200, 200, 200, 200), + (("llm_api_routes_without_mcp",), 200, 200, 403, 403), + (("mcp_inference_routes",), 403, 403, 403, 403), + ], + ids=["llm_api_routes key", "llm only", "mcp only"], +) +@pytest.mark.asyncio +async def test_an_llm_only_credential_keeps_the_llm_api_carve_outs_except_mcp_discovery( + load_user, + fetch_teams, + allowed_routes, + registered_pass_through, + auto_router_session, + mcp_server_listing, + agent_registry, +): + """Operator-registered pass-through endpoints and the auto-router session read are LLM API + surface that only the ``llm_api_routes`` group used to reach, so an LLM-only credential keeps + them, while the MCP server listing that group also carves out stays MCP surface.""" + minted = await mint_proxy_credential("u1", "team-a", allowed_routes=allowed_routes) + assert isinstance(minted, MintedProxyCredential) + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route", + side_effect=lambda route: route.startswith("/my-llm-endpoint"), + ): + assert _route_status(minted, "/my-llm-endpoint/v1/chat", "POST") == registered_pass_through + assert _route_status(minted, "/auto_router/session", "GET") == auto_router_session + assert _route_status(minted, "/v1/mcp/server", "GET") == mcp_server_listing + assert _route_status(minted, "/v1/agents", "GET") == agent_registry 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_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 69eb01cabcd..e4a3935a58d 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -13,6 +13,7 @@ from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, + JWTIssuerConfig, LiteLLMRoutes, LiteLLM_JWTAuth, LiteLLM_ModelTable, @@ -29,9 +30,11 @@ from litellm.proxy._types import ( from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.auth_checks import TeamNotFoundError from litellm.proxy.auth.handle_jwt import ( + ID_JAG_TYP, JWKS_FETCH_ATTEMPTS, STALE_CACHE_KEY_PREFIX, STALE_WRITTEN_AT_CACHE_KEY_PREFIX, + AssertionRejection, HeaderTeam, JWKSUnreachableError, JWTAuthManager, @@ -8303,3 +8306,198 @@ def test_team_passthrough_route_denial_names_the_matched_deny_entry( assert exc_info.value.status_code == 403 assert expected_detail in str(exc_info.value.detail) + + +ASSERTION_ISSUER = "https://idp.example.com/oauth2/default" +ASSERTION_JWKS_URL = f"{ASSERTION_ISSUER}/v1/keys" +GATEWAY_AUDIENCES = ("https://gateway.example.com/mcp", "https://gateway.example.com") +ID_JAG_CLAIMS = ("iss", "sub", "aud", "exp", "iat", "jti", "client_id") + + +def _assertion_signer(): + import json + + import jwt as pyjwt + from cryptography.hazmat.primitives.asymmetric import rsa + + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk: Final = {**json.loads(pyjwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())), "kid": "k1"} + + def sign(claims: Mapping[str, object], typ: str = ID_JAG_TYP) -> str: + return pyjwt.encode(dict(claims), private_key, algorithm="RS256", headers={"kid": "k1", "typ": typ}) + + return jwk, sign + + +def _id_jag_claims(**overrides): + now: Final = int(time.time()) + return { + "iss": ASSERTION_ISSUER, + "sub": "okta-user-1", + "aud": GATEWAY_AUDIENCES[0], + "exp": now + 300, + "iat": now - 5, + "jti": "jti-1", + "client_id": "0oa-requesting-app", + **overrides, + } + + +def _assertion_handler(jwk, litellm_jwtauth, jwks_url=ASSERTION_JWKS_URL): + """A handler whose JWKS for the issuer is already cached, so no HTTP is needed.""" + handler: Final = JWTHandler() + cache: Final = DualCache() + cache.set_cache(f"litellm_jwt_auth_keys_{jwks_url}", [jwk]) + handler.update_environment(prisma_client=None, user_api_key_cache=cache, litellm_jwtauth=litellm_jwtauth) + return handler + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_verifies_against_the_configured_issuer_with_the_callers_audience(): + """The assertion is checked with the issuer's JWKS and claim mappings, but for the audience the + caller names (the gateway's own issuer values), not the API audience the issuer entry carries, + since an ID-JAG is addressed to the authorization server, never to the API.""" + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[ + JWTIssuerConfig( + issuer=ASSERTION_ISSUER, + jwks_url=ASSERTION_JWKS_URL, + audience="api://litellm", + user_id_jwt_field="sub", + ) + ] + ), + ) + claims = _id_jag_claims() + verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(verified, dict) + assert {key: verified[key] for key in claims} == claims + assert handler.get_user_id(token=verified, default_value=None) == "okta-user-1" + second = await handler.decode_trusted_assertion( + sign(_id_jag_claims(aud=GATEWAY_AUDIENCES[1])), GATEWAY_AUDIENCES, ID_JAG_CLAIMS + ) + assert isinstance(second, dict) + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_falls_back_to_the_env_scoped_issuer(monkeypatch): + jwk, sign = _assertion_signer() + monkeypatch.setenv("JWT_ISSUER", ASSERTION_ISSUER) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", ASSERTION_JWKS_URL) + handler = _assertion_handler(jwk, LiteLLM_JWTAuth()) + verified = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(verified, dict) + assert verified["jti"] == "jti-1" + monkeypatch.setenv("JWT_ISSUER", "https://someone-else.example.com") + foreign = await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(foreign, AssertionRejection) + assert "trusts" in foreign.reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fault, reason", + [ + ({"aud": "api://litellm"}, "udience"), + ({"jti": None}, "jti"), + ({"client_id": None}, "client_id"), + ({"exp": -10}, "xpired"), + ({"iat": 600}, "iat"), + ({"iss": "https://untrusted.example.com"}, "trusts"), + ], +) +async def test_decode_trusted_assertion_rejects_a_claim_outside_the_profile(fault: dict[str, object], reason: str): + """The claims are built inside the test, relative to the clock the verifier reads, so a slow + collection cannot turn an intended claim fault into an unrelated expiry (``None`` drops a claim, + an int on ``exp`` or ``iat`` is an offset from now).""" + now: Final = int(time.time()) + offsets: Final = {k: now + v for k, v in fault.items() if k in ("exp", "iat") and isinstance(v, int)} + claims: Final = {k: v for k, v in _id_jag_claims(**offsets).items() if fault.get(k, k) is not None} + claims.update({k: v for k, v in fault.items() if v is not None and k not in offsets}) + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + ) + rejected = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(rejected, AssertionRejection) + assert reason in rejected.reason + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_rejects_a_forged_signature_and_garbage(monkeypatch): + jwk, _ = _assertion_signer() + _, forge = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + ) + forged = await handler.decode_trusted_assertion(forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(forged, AssertionRejection) + assert "verification" in forged.reason + garbage = await handler.decode_trusted_assertion("not.a.jwt", GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert isinstance(garbage, AssertionRejection) + monkeypatch.delenv("JWT_ISSUER", raising=False) + unconfigured = await _assertion_handler(jwk, LiteLLM_JWTAuth()).decode_trusted_assertion( + forge(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS + ) + assert isinstance(unconfigured, AssertionRejection) + assert "trusts" in unconfigured.reason + + +@pytest.mark.asyncio +async def test_decode_trusted_assertion_reports_an_unreachable_jwks_as_an_outage_not_a_bad_assertion(): + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth(issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="x")]), + jwks_url="https://a-different-cached-url.example.com/keys", + ) + with patch.object(handler, "_fetch_jwks_keys", new_callable=AsyncMock, side_effect=JWKSUnreachableError("down")): + with pytest.raises(ProxyException) as outage: + await handler.decode_trusted_assertion(sign(_id_jag_claims()), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert outage.value.code == "503" + + +@pytest.mark.asyncio +async def test_auth_jwt_never_accepts_an_id_jag_as_a_bearer_token(): + """An assertion typed oauth-id-jag+jwt is a grant for the token endpoint, so bearer auth + refuses it before any key lookup, even where it would otherwise verify.""" + jwk, sign = _assertion_signer() + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[ + JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience=GATEWAY_AUDIENCES[0]) + ] + ), + ) + accepted = await handler.auth_jwt(sign(_id_jag_claims(), typ="JWT")) + assert accepted["jti"] == "jti-1" + with pytest.raises(Exception, match="token endpoint"): + await handler.auth_jwt(sign(_id_jag_claims())) +@pytest.mark.asyncio +@pytest.mark.parametrize("verdict", [True, False]) +async def test_decode_trusted_assertion_runs_the_configured_custom_validate_on_the_verified_claims(verdict): + """The operator's ``custom_validate`` gates every JWT the proxy admits, so an assertion redeemed + at the token endpoint goes through it too, after signature and claim verification.""" + jwk, sign = _assertion_signer() + seen: Final = [] + handler = _assertion_handler( + jwk, + LiteLLM_JWTAuth( + issuers=[JWTIssuerConfig(issuer=ASSERTION_ISSUER, jwks_url=ASSERTION_JWKS_URL, audience="api://litellm")] + ), + ) + handler.litellm_jwtauth.custom_validate = lambda claims: seen.append(dict(claims)) or verdict + claims = _id_jag_claims() + verified = await handler.decode_trusted_assertion(sign(claims), GATEWAY_AUDIENCES, ID_JAG_CLAIMS) + assert [{key: call[key] for key in claims} for call in seen] == [claims] + if verdict: + assert isinstance(verified, dict) + else: + assert isinstance(verified, AssertionRejection) + assert "custom_validate" in verified.reason 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/db/test_spend_log_tool_index.py b/tests/unit/proxy/db/test_spend_log_tool_index.py index eaa94c0438f..af3268a7cc9 100644 --- a/tests/unit/proxy/db/test_spend_log_tool_index.py +++ b/tests/unit/proxy/db/test_spend_log_tool_index.py @@ -21,6 +21,8 @@ from litellm.proxy.db.spend_log_tool_index import ( flush_tool_usage_transactions, response_tool_call_names, ) +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -222,6 +224,37 @@ class TestResponseToolCallNames: } assert response_tool_call_names(response) == ("ant_get_weather",) + def test_auto_executed_mcp_tools_count_from_the_final_responses_api_response(self): + final_round = ResponsesAPIResponse.model_validate( + { + "id": "resp_final", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + } + ) + logged = LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + final_round, + mcp_tools_fetched=[{"name": "demo-echo"}, {"name": "demo-add"}], + tool_results=[ + {"tool_call_id": "call_1", "result": "x", "name": "demo-echo"}, + {"tool_call_id": "call_2", "result": "3", "name": "demo-add"}, + ], + ) + + assert response_tool_call_names(logged) == ("demo-echo", "demo-add") + assert response_tool_call_names(logged.model_dump()) == ("demo-echo", "demo-add") + def _transaction( request_id: str, 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/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 33904637366..c6369a1f3d7 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -20,7 +20,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_start_time_cache_key, resolve_model_budget, ) -from litellm.types.caching import RedisPipelineIncrementOperation async def _reset_at_of_open_window( @@ -500,16 +499,15 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config @pytest.mark.asyncio -async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): +async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queues_nothing(): dual_cache = DualCache() limiter = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) spend_key = "virtual_key_spend:test-key:gpt-4:1d" await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.01, ttl=86400) - assert limiter.redis_increment_operation_queue == [ - RedisPipelineIncrementOperation(key=spend_key, increment_value=0.01, ttl=86400) - ] + assert await dual_cache.async_get_cache(spend_key) == 0.01 + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py index ff4fbbfb695..97c36fd5792 100644 --- a/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py +++ b/tests/unit/proxy/management_endpoints/test_id_jag_assertion_capture.py @@ -49,8 +49,8 @@ class TestActiveSSOProviderMirrorsTheCallback: assert active_sso_provider() is ActiveSSOProvider.none def test_google_outranks_a_co_configured_generic_client(self, monkeypatch): - """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and - no assertion is captured; reporting generic would clear a gap that is still open.""" + """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and the + assertion the refresher later renews was minted by the Google client, not the generic one.""" monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") assert active_sso_provider() is ActiveSSOProvider.google @@ -67,51 +67,41 @@ class TestActiveSSOProviderMirrorsTheCallback: class TestIdJagAssertionCaptureGap: - def test_generic_oidc_has_no_gap(self, monkeypatch): - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + @pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + def test_every_oidc_login_captures_so_has_no_gap(self, monkeypatch, env_var): + """Google, Microsoft and generic OIDC all read the id_token off the token response, so a + deployment on any of them must not be told to switch providers.""" + monkeypatch.setenv(env_var, "configured") assert id_jag_assertion_capture_gap() is None - @pytest.mark.parametrize( - "env_var, provider_label", - [ - ("GOOGLE_CLIENT_ID", "google"), - ("MICROSOFT_CLIENT_ID", "microsoft"), - ("SAML_IDP_METADATA_URL", "saml"), - ], - ) - def test_non_capturing_provider_is_named_with_the_remedy(self, monkeypatch, env_var, provider_label): - monkeypatch.setenv(env_var, "configured") - gap = id_jag_assertion_capture_gap() - assert gap is not None - assert provider_label in gap - assert "GENERIC_CLIENT_ID" in gap + def test_saml_has_no_gap(self, monkeypatch: pytest.MonkeyPatch): + """The SAML callback retains the signed assertion as the saml2 subject token, so a SAML + deployment is a capturing one and must not be told to switch to an OIDC provider.""" + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert id_jag_assertion_capture_gap() is None def test_no_sso_configured_reports_a_gap(self): gap = id_jag_assertion_capture_gap() assert gap is not None assert "no SSO provider is configured" in gap - def test_google_beside_generic_still_reports_a_gap(self, monkeypatch): - """The precedence trap in operator terms: adding a generic client id without removing - GOOGLE_CLIENT_ID does not fix the deployment, so the gap must not clear.""" - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") - gap = id_jag_assertion_capture_gap() - assert gap is not None - assert "google" in gap + def test_saml_beside_an_oidc_client_has_no_gap(self, monkeypatch): + """The callback dispatches to the OIDC client before SAML, so that is the login that runs + and it captures; reporting SAML's gap would tell the operator to fix a path nobody takes.""" + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert id_jag_assertion_capture_gap() is None class TestIdJagAssertionCaptureGapAtStartup: def test_no_provider_at_startup_is_not_yet_a_gap(self): assert id_jag_assertion_capture_gap_at_startup() is None - def test_google_provider_at_startup_reports_the_capture_gap(self, monkeypatch): - monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") - startup_gap = id_jag_assertion_capture_gap_at_startup() - callback_gap = id_jag_assertion_capture_gap() - assert startup_gap is not None - assert startup_gap == callback_gap - - def test_generic_provider_at_startup_has_no_gap(self, monkeypatch): - monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + def test_saml_at_startup_has_no_gap(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert id_jag_assertion_capture_gap_at_startup() is None + + @pytest.mark.parametrize("env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + def test_an_oidc_provider_at_startup_has_no_gap(self, monkeypatch, env_var): + monkeypatch.setenv(env_var, "configured") assert id_jag_assertion_capture_gap_at_startup() is None diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 47aea89843f..0d59c655835 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4458,9 +4458,9 @@ class TestAddMCPServerAtomicity: class TestIdJagRegistrationWarnsAboutTheSSOGap: """An `oauth2_id_jag` server only ever works when the login path captures an IdP identity - assertion, and only the generic OIDC arm does. Registering one under Google or Microsoft - succeeds and then fails for every user on every call, so the mismatch has to be said at - registration time, while the admin is still looking at the configuration.""" + assertion, and only the OIDC logins (Google, Microsoft, generic) do. Registering one under SAML + or with no SSO succeeds and then fails for every user on every call, so the mismatch has to be + said at registration time, while the admin is still looking at the configuration.""" @staticmethod def _clear_sso_env(monkeypatch): @@ -4530,9 +4530,6 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"GOOGLE_CLIENT_ID": "cid"}, "google"), - ({"MICROSOFT_CLIENT_ID": "cid"}, "microsoft"), - ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), ({}, "no SSO provider is configured"), ], ) @@ -4546,8 +4543,9 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: assert "ema-1" in str(warnings[0]) @pytest.mark.asyncio - async def test_create_is_silent_under_generic_oidc(self, monkeypatch, caplog): - await self._run_create(monkeypatch, {"GENERIC_CLIENT_ID": "cid"}, MCPAuth.oauth2_id_jag, caplog) + @pytest.mark.parametrize("client_id_env", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) + async def test_create_is_silent_under_oidc_logins(self, monkeypatch, caplog, client_id_env): + await self._run_create(monkeypatch, {client_id_env: "cid"}, MCPAuth.oauth2_id_jag, caplog) assert self._id_jag_warnings(caplog) == [] @pytest.mark.asyncio @@ -4565,7 +4563,8 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: ) self._clear_sso_env(monkeypatch) - monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(var, raising=False) mock_manager = MagicMock() mock_manager.update_server = AsyncMock() @@ -4602,7 +4601,7 @@ class TestIdJagRegistrationWarnsAboutTheSSOGap: warnings = self._id_jag_warnings(caplog) assert len(warnings) == 1 - assert "google" in str(warnings[0]) + assert "no SSO provider is configured" in str(warnings[0]) class TestHealthCheckServers: 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_saml_sso.py b/tests/unit/proxy/management_endpoints/test_saml_sso.py index 01f8fc929a3..181c5e65dae 100644 --- a/tests/unit/proxy/management_endpoints/test_saml_sso.py +++ b/tests/unit/proxy/management_endpoints/test_saml_sso.py @@ -9,15 +9,17 @@ makes a test fail. import base64 import datetime +import logging +import re import time from http.cookies import SimpleCookie import pytest from fastapi import HTTPException, Request -pytest.importorskip( - "onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests" -) +pytest.importorskip("onelogin", reason="python3-saml (saml extra) is required for SAML SSO tests") + +from typing import cast from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization @@ -26,11 +28,10 @@ from cryptography.x509.oid import NameOID from onelogin.saml2.utils import OneLogin_Saml2_Utils from starlette.datastructures import URL -from typing import cast - from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import SSOIdentityAssertion from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.sso.saml_sso import ( _SAML_AUTHN_REQUEST_CACHE_PREFIX, @@ -40,15 +41,17 @@ from litellm.proxy.management_endpoints.sso.saml_sso import ( _SAML_REPLAY_GUARD_MAX_TTL_SECONDS, SAMLAuthHandler, ) +from litellm.proxy.management_endpoints.types import CustomOpenID -def _shared_cache(store=None): +def _shared_cache(store: InMemoryCache | None = None) -> DualCache: """A DualCache whose replay guard is backed by a shared, atomic store. An InMemoryCache instance stands in for Redis; passing the same instance to two DualCaches simulates two workers sharing one atomic backend.""" return DualCache(redis_cache=cast(RedisCache, store or InMemoryCache())) + IDP_ENTITY = "https://idp.example.com/metadata" SP_ENTITY = "https://proxy.example.com/sso/saml/metadata" ACS = "https://proxy.example.com/sso/saml/callback" @@ -56,7 +59,7 @@ SSO_URL = "https://idp.example.com/sso" PROXY_BASE_URL = "https://proxy.example.com" -def _make_idp_keypair(): +def _make_idp_keypair() -> tuple[str, str]: key = rsa.generate_private_key(public_exponent=65537, key_size=2048) name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "idp.example.com")]) cert = ( @@ -79,9 +82,7 @@ def _make_idp_keypair(): def _idp_metadata_xml(cert_pem): - cert_body = "".join( - line for line in cert_pem.splitlines() if "CERTIFICATE" not in line - ) + cert_body = "".join(line for line in cert_pem.splitlines() if "CERTIFICATE" not in line) return ( '' f'' @@ -101,17 +102,17 @@ def _saml_time(delta_seconds): def _build_signed_response( - key_pem, - cert_pem, + key_pem: str, + cert_pem: str, *, - in_response_to=None, - response_level_in_response_to=True, - email="alice@example.com", - attributes=None, - not_before_delta=-60, - not_on_or_after_delta=300, - sign=True, -): + in_response_to: str | None = None, + response_level_in_response_to: bool = True, + email: str = "alice@example.com", + attributes: dict[str, list[str]] | None = None, + not_before_delta: int = -60, + not_on_or_after_delta: int = 300, + sign: bool = True, +) -> str: if attributes is None: attributes = { "email": [email], @@ -157,9 +158,7 @@ def _build_signed_response( if sign: signed = OneLogin_Saml2_Utils.add_sign(assertion, key_pem, cert_pem) - assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace( - '', "" - ) + assertion = (signed.decode() if isinstance(signed, bytes) else signed).replace('', "") return ( '' @@ -174,11 +173,11 @@ def _build_signed_response( ) -def _b64(xml): +def _b64(xml: str) -> str: return base64.b64encode(xml.encode()).decode() -def _fake_request(cookies=None): +def _fake_request(cookies: dict[str, str] | None = None): return type( "Req", (), @@ -190,14 +189,18 @@ def _fake_request(cookies=None): )() +async def _acs_login( + b64: str, cache: DualCache, cookies: dict[str, str] | None = None +) -> tuple[CustomOpenID, SSOIdentityAssertion | None]: + return await SAMLAuthHandler.handle_acs(_fake_request(cookies), cache, {"SAMLResponse": b64}) + + async def _acs(b64, cache, cookies=None): - return await SAMLAuthHandler.handle_acs( - _fake_request(cookies), cache, {"SAMLResponse": b64} - ) + return (await _acs_login(b64, cache, cookies))[0] @pytest.fixture -def saml_env(monkeypatch): +def saml_env(monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]: key_pem, cert_pem = _make_idp_keypair() monkeypatch.setenv("SAML_IDP_METADATA_XML", _idp_metadata_xml(cert_pem)) monkeypatch.setenv("SAML_SP_ENTITY_ID", SP_ENTITY) @@ -214,7 +217,7 @@ def saml_env(monkeypatch): @pytest.fixture -def saml_env_idp_initiated(saml_env, monkeypatch): +def saml_env_idp_initiated(saml_env: tuple[str, str], monkeypatch: pytest.MonkeyPatch) -> tuple[str, str]: monkeypatch.setenv("SAML_ALLOW_UNSOLICITED", "true") return saml_env @@ -234,6 +237,89 @@ async def test_valid_idp_initiated_login_maps_assertion_to_user(saml_env_idp_ini assert result.provider == "saml" +@pytest.mark.asyncio +async def test_valid_login_retains_the_signed_assertion_as_a_saml2_subject_token( + saml_env_idp_initiated: tuple[str, str], +): + """The retained material is the inner , still carrying a verifiable signature, with + the assertion's own NotOnOrAfter as its expiry and no refresh token, since SAML has none.""" + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, not_on_or_after_delta=300) + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert assertion is not None + assert assertion.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" + assert assertion.refresh_token is None + assert assertion.issuer == IDP_ENTITY + retained = _retained_bytes(assertion.id_token.get_secret_value()) + assert retained.startswith(b" str: + match = re.search(r']*NotOnOrAfter="([^"]+)"', response_xml) + assert match is not None + return match.group(1) + + +def _signature_valid(xml: bytes, cert_pem: str) -> bool: + return bool(OneLogin_Saml2_Utils.validate_sign(xml, cert_pem, raise_exceptions=True, xpath="//ds:Signature")) + + +def _retained_bytes(encoded: str) -> bytes: + assert set(encoded) <= set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"), encoded[:40] + return base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)) + + +@pytest.mark.asyncio +async def test_retained_assertion_keeps_non_ascii_subject_bytes_verbatim(saml_env_idp_initiated: tuple[str, str]): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, email="zo\u00eb@example.com") + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert assertion is not None + retained = _retained_bytes(assertion.id_token.get_secret_value()) + assert "zo\u00eb@example.com".encode() in retained + assert b"&#" not in retained + assert _signature_valid(retained, cert_pem) + + +class _UnparseableResponseAuth: + def get_last_response_xml(self) -> bytes: + return b" int | None: + return None + + +def test_assertion_retention_failure_does_not_fail_the_validated_login(caplog: pytest.LogCaptureFixture): + """Retention is optional egress material, so a response the retention step cannot parse logs and yields + no assertion instead of turning an IdP-validated login into a 500.""" + with caplog.at_level(logging.WARNING): + retained = SAMLAuthHandler._retained_assertion(_UnparseableResponseAuth()) + + assert retained is None + assert "could not be retained for oauth2_id_jag" in caplog.text + + +@pytest.mark.asyncio +async def test_retained_assertion_is_the_one_the_user_was_mapped_from(saml_env_idp_initiated: tuple[str, str]): + key_pem, cert_pem = saml_env_idp_initiated + resp = _build_signed_response(key_pem, cert_pem, email="bob@example.com") + + user, assertion = await _acs_login(_b64(resp), _shared_cache()) + + assert user.email == "bob@example.com" + assert assertion is not None + assert b"bob@example.com" in _retained_bytes(assertion.id_token.get_secret_value()) + + @pytest.mark.asyncio async def test_tampered_assertion_is_rejected(saml_env): key_pem, cert_pem = saml_env @@ -269,9 +355,7 @@ async def test_signature_from_untrusted_key_is_rejected(saml_env): @pytest.mark.asyncio async def test_expired_assertion_is_rejected(saml_env): key_pem, cert_pem = saml_env - resp = _build_signed_response( - key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600 - ) + resp = _build_signed_response(key_pem, cert_pem, not_before_delta=-7200, not_on_or_after_delta=-3600) with pytest.raises(HTTPException) as exc: await _acs(_b64(resp), DualCache()) @@ -293,9 +377,7 @@ async def test_sp_initiated_known_request_succeeds_once_then_replay_rejected(sam key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) cookies = {_SAML_AUTHN_STATE_COOKIE: request_id} @@ -312,9 +394,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response(key_pem, cert_pem, in_response_to=request_id) with pytest.raises(HTTPException) as exc: @@ -322,9 +402,7 @@ async def test_sp_initiated_response_not_bound_to_browser_is_rejected(saml_env): assert exc.value.status_code == 401 with pytest.raises(HTTPException) as exc: - await _acs( - _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"} - ) + await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: "_attacker_request"}) assert exc.value.status_code == 401 @@ -338,9 +416,7 @@ async def test_subjectconfirmation_only_in_response_to_without_cookie_is_rejecte key_pem, cert_pem = saml_env_idp_initiated cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response( key_pem, cert_pem, @@ -358,9 +434,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml key_pem, cert_pem = saml_env cache = DualCache() request_id = "_authn_req_known" - cache.set_cache( - key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600 - ) + cache.set_cache(key=f"{_SAML_AUTHN_REQUEST_CACHE_PREFIX}:{request_id}", value="1", ttl=600) resp = _build_signed_response( key_pem, cert_pem, @@ -368,9 +442,7 @@ async def test_subjectconfirmation_only_in_response_to_with_cookie_succeeds(saml response_level_in_response_to=False, ) - result = await _acs( - _b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id} - ) + result = await _acs(_b64(resp), cache, cookies={_SAML_AUTHN_STATE_COOKIE: request_id}) assert result.email == "alice@example.com" @@ -415,9 +487,7 @@ async def test_assertion_without_id_is_rejected(saml_env_idp_initiated): return None with pytest.raises(HTTPException) as exc: - await SAMLAuthHandler._enforce_response_binding( - _AuthNoAssertionId(), _shared_cache(), None - ) + await SAMLAuthHandler._enforce_response_binding(_AuthNoAssertionId(), _shared_cache(), None) assert exc.value.status_code == 401 assert "ID" in exc.value.detail @@ -450,9 +520,7 @@ async def test_invalid_email_in_assertion_is_rejected_cleanly(saml_env_idp_initi @pytest.mark.asyncio -async def test_email_less_assertion_rejected_when_domain_restriction_configured( - saml_env_idp_initiated, monkeypatch -): +async def test_email_less_assertion_rejected_when_domain_restriction_configured(saml_env_idp_initiated, monkeypatch): key_pem, cert_pem = saml_env_idp_initiated monkeypatch.setenv("ALLOWED_EMAIL_DOMAINS", "example.com") resp = _build_signed_response( @@ -558,11 +626,7 @@ async def test_build_login_redirect_targets_idp_and_caches_request_id(saml_env): location = redirect.headers["location"] assert location.startswith(SSO_URL) assert "SAMLRequest=" in location - cached = [ - k - for k in cache.in_memory_cache.cache_dict - if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX) - ] + cached = [k for k in cache.in_memory_cache.cache_dict if k.startswith(_SAML_AUTHN_REQUEST_CACHE_PREFIX)] assert len(cached) == 1 request_id = cached[0].split(":", 1)[1] @@ -761,9 +825,7 @@ class TestSAMLAuthnCookieSecureFlag: {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, ) cache = DualCache() - request = _fake_request_with_scheme( - "http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5" - ) + request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5") redirect = await SAMLAuthHandler.build_login_redirect(request, cache) cookie = redirect.headers["set-cookie"] assert "Secure" in cookie @@ -773,9 +835,32 @@ class TestSAMLAuthnCookieSecureFlag: monkeypatch.delenv("PROXY_BASE_URL", raising=False) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) cache = DualCache() - request = _fake_request_with_scheme( - "http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5" - ) + request = _fake_request_with_scheme("http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5") redirect = await SAMLAuthHandler.build_login_redirect(request, cache) cookie = redirect.headers["set-cookie"] assert "Secure" not in cookie + + +def test_an_encrypted_assertion_is_not_retained_and_the_login_log_says_why(caplog: pytest.LogCaptureFixture): + """An IdP that encrypts its assertion still signs the user in, but the ID-JAG arm has nothing to + present. The login log has to name that, otherwise every later 412 looks like a bug re-login cannot fix.""" + + class _Encrypted: + def get_last_response_xml(self) -> bytes: + return ( + b'' + b'' + b"" + ) + + def get_last_assertion_not_on_or_after(self) -> int | None: + return None + + with caplog.at_level(logging.WARNING, logger="litellm.proxy"): + retained = SAMLAuthHandler._retained_assertion(_Encrypted()) # pyright: ignore[reportArgumentType] # stub auth # pyright: ignore[reportPrivateUsage] # the retention seam under test + + assert retained is None + assert any( + "1 saml:EncryptedAssertion" in r.getMessage() and "oauth2_id_jag" in r.getMessage() for r in caplog.records + ) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index f69c069cfea..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 @@ -10,13 +10,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from pydantic import SecretStr import respx 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 @@ -25,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, @@ -243,17 +246,24 @@ def test_get_microsoft_callback_response(stubbed_graph_api): "surname": "User", } + id_token = _ema_id_token(sub="msft123") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = "msft-refresh-token" + return mock_response + with patch.dict( os.environ, {"MICROSOFT_CLIENT_SECRET": "mock_secret", "MICROSOFT_TENANT": "mock_tenant"}, ): - mock_verify = AsyncMock(return_value=mock_response) - with patch( + with patch( # test-quality-ok: stubs the upstream Microsoft IdP boundary, where the token response is read "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", - new=mock_verify, + autospec=True, + side_effect=fake_verify, ): # Act - result = asyncio.run( + result, sso_assertion = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -269,6 +279,11 @@ def test_get_microsoft_callback_response(stubbed_graph_api): assert result.id == "msft123" assert result.first_name == "Microsoft" assert result.last_name == "User" + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is not None + assert sso_assertion.refresh_token.get_secret_value() == "msft-refresh-token" + assert sso_assertion.issuer == "https://idp.example.com" GRAPH_ERROR_BODIES: Final = ( @@ -317,7 +332,7 @@ def test_get_microsoft_callback_response_signs_in_a_profile_with_an_empty_error_ "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", new=AsyncMock(return_value={**profile, "error": empty_error}), ): - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -342,7 +357,7 @@ def test_get_microsoft_callback_response_raw_sso_response_keeps_the_graph_error( "fastapi_sso.sso.microsoft.MicrosoftSSO.verify_and_process", new=AsyncMock(return_value=dict(graph_error)), ): - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -376,7 +391,7 @@ def test_get_microsoft_callback_response_raw_sso_response(stubbed_graph_api): new=mock_verify, ): # Act - result = asyncio.run( + result, _ = asyncio.run( MicrosoftSSOHandler.get_microsoft_callback_response( request=mock_request, microsoft_client_id="mock_client_id", @@ -405,13 +420,19 @@ def test_get_google_callback_response(): "family_name": "User", } + id_token = _ema_id_token(sub="google123") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = None + return mock_response + with patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "mock_secret"}): - mock_verify = AsyncMock(return_value=mock_response) - with patch( - "fastapi_sso.sso.google.GoogleSSO.verify_and_process", new=mock_verify + with patch( # test-quality-ok: stubs the upstream Google IdP boundary, where the token response is read + "fastapi_sso.sso.google.GoogleSSO.verify_and_process", autospec=True, side_effect=fake_verify ): # Act - result = asyncio.run( + result, sso_assertion = asyncio.run( GoogleSSOHandler.get_google_callback_response( request=mock_request, google_client_id="mock_client_id", @@ -426,6 +447,9 @@ def test_get_google_callback_response(): assert result.get("sub") == "google123" assert result.get("given_name") == "Google" assert result.get("family_name") == "User" + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is None @pytest.mark.asyncio @@ -1403,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. @@ -4018,7 +4118,13 @@ class TestCLIKeyRegenerationFlow: cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456" # Mock the CLI callback and required proxy server components + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + mock_result = {"user_id": "test-user", "email": "test@example.com"} + google_assertion = assertion_from_sso_login(_ema_id_token(sub="test-user"), None) + assert google_assertion is not None with ( patch( @@ -4032,7 +4138,7 @@ class TestCLIKeyRegenerationFlow: patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True), patch( "litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response", - return_value=mock_result, + return_value=(mock_result, google_assertion), ), ): mock_cli_callback.return_value = MagicMock() @@ -4046,7 +4152,7 @@ class TestCLIKeyRegenerationFlow: prefill_user_code=None, result=mock_result, received_response=None, - sso_assertion=None, + sso_assertion=google_assertion, ) @pytest.mark.asyncio @@ -4074,7 +4180,7 @@ class TestCLIKeyRegenerationFlow: patch.dict(os.environ, {"GOOGLE_CLIENT_ID": "test-google-id"}, clear=True), patch( "litellm.proxy.management_endpoints.ui_sso.GoogleSSOHandler.get_google_callback_response", - return_value=mock_result, + return_value=(mock_result, None), ), ): mock_cli_callback.return_value = MagicMock() @@ -8836,7 +8942,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): ) async def fake_microsoft_callback(**kwargs): - return parsed_openid + return parsed_openid, None with ( patch.dict( @@ -8886,7 +8992,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) return parsed, {"sub": "user_123"}, {"scope": "openid"}, None async def fake_google(**kwargs): - return parsed + return parsed, None stack = [ patch.dict(os.environ, {**provider_env, "ENABLE_SSO_DEBUG": "true"}, clear=False), @@ -8925,14 +9031,14 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) @pytest.mark.asyncio -async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): +async def test_debug_page_logs_no_capture_gap_for_a_google_login(caplog): + """Google reads the id_token off its token response like generic OIDC does, so the debug page + under GOOGLE_CLIENT_ID with an oauth2_id_jag server registered has no gap to log and nothing + to render about it.""" with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) - warnings = _id_jag_gap_warnings(caplog) - assert len(warnings) == 1 - assert "google" in warnings[0] - assert "GENERIC_CLIENT_ID" in warnings[0] + assert _id_jag_gap_warnings(caplog) == [] assert "id_jag" not in body assert "GENERIC_CLIENT_ID" not in body @@ -8977,7 +9083,8 @@ async def test_debug_page_survives_a_store_outage(monkeypatch, caplog): not the page.""" from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap - monkeypatch.setenv("GOOGLE_CLIENT_ID", _GOOGLE_DEBUG_CLIENT_ID) + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(var, raising=False) retention_check = AsyncMock(side_effect=Exception("db down")) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None @@ -9103,7 +9210,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): async def _fake_handle_acs(**kwargs): call_order.append("validate") - return CustomOpenID( + return (CustomOpenID( id="dana@litellm.ai", email="dana@litellm.ai", first_name=None, @@ -9113,7 +9220,7 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): provider="saml", team_ids=[], user_role=None, - ) + ), None) async def _fake_count_billable_users(): call_order.append("count") @@ -9635,6 +9742,233 @@ async def test_redirect_from_openid_persists_assertion_under_canonical_user_id() assert response is not None +@pytest.mark.asyncio +async def test_redirect_from_openid_persists_the_login_exchanged_assertion(): + """A server with ``id_jag_subject_source: refresh_token`` needs the refresh token the login + assertion was exchanged for, so the funnel persists what the login-time exchange returns, not + the raw capture.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + + assertion = assertion_from_sso_login(_ema_id_token(), "rt_1") + assert assertion is not None + mock_request = MagicMock(spec=Request) + mock_request.scope = {} + mock_request.base_url = "http://localhost:4000/" + mock_request.cookies = {} + + retain_mock = AsyncMock() + exchanged = assertion.model_copy(update={"exchanged_refresh_token": SecretStr("rt-from-login")}) + exchange_mock = AsyncMock(return_value=exchanged) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.user_custom_sso", None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + AsyncMock( + return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"} + ), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.check_and_update_if_proxy_admin_id", + AsyncMock(return_value="internal_user"), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + retain_mock, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.exchange_login_assertion_if_configured", + exchange_mock, + ), + ): + response = await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=CustomOpenID( + id="raw-idp-subject", + email="u@example.com", + first_name="U", + last_name="Ser", + display_name="U Ser", + provider="generic", + team_ids=[], + user_role=None, + ), + request=mock_request, + received_response=None, + generic_client_id="cid", + ui_access_mode=None, + access_token_payload=None, + jwt_handler=None, + sso_assertion=assertion, + ) + exchange_mock.assert_awaited_once_with(assertion) + retain_mock.assert_awaited_once() + assert retain_mock.await_args.kwargs["assertion"] is exchanged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_env, sso_class, payload, needs_graph", + [ + ( + {"MICROSOFT_CLIENT_ID": "ms-cid", "MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"}, + "fastapi_sso.sso.microsoft.MicrosoftSSO", + {"mail": "u@example.com", "displayName": "U Ser", "id": "idp-subject"}, + True, + ), + ( + {"GOOGLE_CLIENT_ID": "google-cid", "GOOGLE_CLIENT_SECRET": "google-secret"}, + "fastapi_sso.sso.google.GoogleSSO", + {"email": "u@example.com", "name": "U Ser", "sub": "idp-subject"}, + False, + ), + ], +) +async def test_ui_login_through_microsoft_or_google_hands_its_assertion_to_the_redirect( + request, provider_env, sso_class, payload, needs_graph +): + """The UI funnel used to capture an assertion only on the generic branch, so an ``oauth2_id_jag`` + server had nothing to spend for a deployment signing in through Entra or Google. The real + Microsoft and Google handlers run here, only the IdP token exchange is stubbed, so the id_token + and refresh token the IdP returned must reach the redirect that persists them.""" + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, auth_callback + + if needs_graph: + request.getfixturevalue("stubbed_graph_api") + id_token = _ema_id_token(sub="idp-subject") + + async def fake_verify(self, *args, **kwargs): + self._id_token = id_token + self._refresh_token = "idp-refresh" + return payload + + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"code": "auth-code"} + mock_request.cookies = {} + mock_request.scope = {} + mock_request.base_url = "http://localhost:4000/" + redirect_mock = AsyncMock(return_value=MagicMock()) + + with ( + patch.dict(os.environ, provider_env, clear=True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch( # test-quality-ok: stubs the upstream IdP token exchange, where the token response is read + f"{sso_class}.verify_and_process", autospec=True, side_effect=fake_verify + ), + patch.object( # test-quality-ok: the DB and key-mint boundary after the login, asserted on its inputs + SSOAuthenticationHandler, "get_redirect_response_from_openid", redirect_mock + ), + ): + await auth_callback(request=mock_request, state=None) + + redirect_mock.assert_awaited_once() + assert redirect_mock.await_args.kwargs["result"] is not None + sso_assertion = redirect_mock.await_args.kwargs["sso_assertion"] + assert sso_assertion is not None + assert sso_assertion.id_token.get_secret_value() == id_token + assert sso_assertion.refresh_token is not None + assert sso_assertion.refresh_token.get_secret_value() == "idp-refresh" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retention_enabled, expected_params", + [ + (AsyncMock(return_value=True), {"access_type": "offline", "prompt": "consent"}), + (AsyncMock(return_value=False), None), + (AsyncMock(side_effect=RuntimeError("store down")), None), + ], +) +async def test_google_login_asks_for_a_refresh_token_only_while_an_id_jag_server_exists( + retention_enabled, expected_params +): + """Google issues a refresh token only when the authorization request carries ``access_type=offline``, + and the refresher cannot renew a Google assertion without one. The login asks only while an + ``oauth2_id_jag`` server exists, and a store outage must not fail the login or add the consent.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + seen: dict[str, object] = {} + + async def fake_redirect(self, *, redirect_uri=None, params=None, state=None): + seen["params"] = params + return RedirectResponse(url="https://accounts.google.com/o/oauth2/v2/auth") + + with ( + patch.dict(os.environ, {"GOOGLE_CLIENT_SECRET": "google-secret"}), + patch( # test-quality-ok: stubs the upstream Google authorization boundary, where the request leaves + "fastapi_sso.sso.google.GoogleSSO.get_login_redirect", autospec=True, side_effect=fake_redirect + ), + ): + response = await SSOAuthenticationHandler.get_sso_login_redirect( + redirect_url="http://localhost:4000/sso/callback", + google_client_id="google-cid", + retention_enabled=retention_enabled, + ) + + assert response is not None + assert seen["params"] == expected_params + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retention_enabled, asks_offline_access", + [ + (AsyncMock(return_value=True), True), + (AsyncMock(return_value=False), False), + (AsyncMock(side_effect=RuntimeError("store down")), False), + ], +) +async def test_microsoft_login_asks_for_offline_access_only_while_an_id_jag_server_exists( + retention_enabled, asks_offline_access +): + """Entra issues a refresh token only when ``offline_access`` is in the requested scope, and the + refresher cannot renew an Entra assertion without one. The default sign-in scopes stay in place so + the login keeps resolving the same profile.""" + from fastapi.responses import RedirectResponse + from fastapi_sso.sso.microsoft import MicrosoftSSO + + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + seen: dict[str, object] = {} + + async def fake_redirect(self, *, redirect_uri=None, params=None, state=None): + seen["scope"] = list(self._scope) + return RedirectResponse(url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize") + + with ( + patch.dict(os.environ, {"MICROSOFT_CLIENT_SECRET": "ms-secret", "MICROSOFT_TENANT": "tenant"}), + patch( # test-quality-ok: stubs the upstream Entra authorization boundary, where the request leaves + "fastapi_sso.sso.microsoft.MicrosoftSSO.get_login_redirect", autospec=True, side_effect=fake_redirect + ), + ): + response = await SSOAuthenticationHandler.get_sso_login_redirect( + redirect_url="http://localhost:4000/sso/callback", + microsoft_client_id="ms-cid", + retention_enabled=retention_enabled, + ) + + assert response is not None + assert all(default in seen["scope"] for default in MicrosoftSSO.scope), seen["scope"] + assert ("offline_access" in seen["scope"]) is asks_offline_access, seen["scope"] + + @pytest.mark.asyncio async def test_cli_completion_persists_assertion_under_db_user_id(): """The CLI funnel persists the captured assertion under the DB-resolved user_id.""" @@ -9710,8 +10044,6 @@ def _id_jag_gap_warnings(caplog) -> list[str]: @pytest.mark.parametrize( "provider_env, expected_fragment", [ - ({"GOOGLE_CLIENT_ID": "cid"}, "google"), - ({"MICROSOFT_CLIENT_ID": "cid", "MICROSOFT_TENANT": "t"}, "microsoft"), ({}, "no SSO provider is configured"), ], ) @@ -9738,16 +10070,17 @@ async def test_uncaptured_assertion_warns_when_an_id_jag_server_is_registered( @pytest.mark.asyncio -async def test_generic_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog): - """Generic OIDC has a capture path, so there is no configuration gap to report; the login - still handed the id_jag arm nothing, and that must not pass silently.""" +@pytest.mark.parametrize("provider_env_var", ["GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID"]) +async def test_a_capturing_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog, provider_env_var): + """Google, Microsoft and generic OIDC all have a capture path, so there is no configuration gap + to report; the login still handed the id_jag arm nothing, and that must not pass silently.""" from litellm.proxy.management_endpoints.ui_sso import ( warn_if_id_jag_assertion_uncaptured, ) - for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "SAML_IDP_METADATA_URL"): + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("GENERIC_CLIENT_ID", "cid") + monkeypatch.setenv(provider_env_var, "cid") with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) @@ -9875,7 +10208,7 @@ async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplo warnings = _id_jag_gap_warnings(caplog) assert len(warnings) == 1 - assert "google" in str(warnings[0]) + assert "no usable id_token" in str(warnings[0]) @pytest.mark.asyncio @@ -9939,7 +10272,7 @@ async def test_cli_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): warnings = _id_jag_gap_warnings(caplog) assert len(warnings) == 1 - assert "microsoft" in str(warnings[0]) + assert "no usable id_token" in str(warnings[0]) def _cli_callback_kwargs(flow): 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/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 4c8265b82c3..2d86ff6a7c2 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -14,10 +14,15 @@ Routes covered: from __future__ import annotations import asyncio +import contextlib import json +from collections.abc import Callable +from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from fastapi.testclient import TestClient from .conftest import VOLATILE_KEYS, normalize @@ -944,6 +949,228 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() +def _delete_callback_roundtrip( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, + litellm_settings: dict[str, object], + callback_name: str, + owned_by_config: frozenset[str] = frozenset(), +) -> tuple[httpx.Response, dict[str, object]]: + """POST /config/callback/delete as admin against a stubbed config, and + return (response, saved litellm_settings) so each case pins the full body + plus the exact config persisted.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": litellm_settings}) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = owned_by_config.__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": callback_name}) + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"]["litellm_settings"] + return response, saved + + +def test_config_callback_delete_from_failure_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.failure_callback is + listed by GET /get/config/callbacks, so the delete route must remove it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["slack"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["slack"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_callbacks_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.callbacks is deleted + while unrelated lists pass through untouched.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["langfuse"], "callbacks": ["datadog"], "drop_params": True}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["langfuse"], "callbacks": [], "drop_params": True} + assert response.json()["remaining_callbacks"] == [] + + +def test_config_callback_delete_from_multiple_lists( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name present in both success_callback and callbacks is removed from + both, and remaining_callbacks reports the leftovers of the lists it left.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + { + "success_callback": ["datadog", "slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse", "datadog"], + }, + "datadog", + ) + assert response.status_code == 200 + assert saved == { + "success_callback": ["slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse"], + } + assert response.json()["remaining_callbacks"] == ["slack", "langfuse"] + + +def test_config_callback_delete_skips_lists_the_config_file_owns( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name in a database-owned list and a config-file-owned list is removed + from the database-owned one only, so the save is not refused.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog", "slack"], "failure_callback": ["datadog"]}, + "datadog", + owned_by_config=frozenset({"failure_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["slack"], "failure_callback": ["datadog"]} + assert response.json()["remaining_callbacks"] == ["slack"] + + +def test_config_callback_delete_removes_database_copy_when_config_file_owns_success_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name the config file sets in success_callback is left there while its + database-owned failure_callback copy is removed.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + owned_by_config=frozenset({"success_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["datadog"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_string_valued_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A list written as a single string is matched and saved back as an empty list.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"failure_callback": "datadog"}, + "datadog", + ) + assert response.status_code == 200 + assert saved == {"failure_callback": []} + + +def test_config_callback_delete_case_insensitive( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A mixed-case entry is deleted by the name GET /get/config/callbacks shows for it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"callbacks": ["DataDog", "otel"]}, + "DataDog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"callbacks": ["otel"]} + + +def test_config_callback_delete_saves_against_the_baseline_it_read( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """save_config gets the object get_config returned, so it diffs against the + settings as they were read and writes only the edited callback list.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + read_time: Final = {"litellm_settings": {"failure_callback": ["datadog", "sentry"], "max_budget": 10}} + loaded: Final = ps._ConfigWithBaseline(read_time) + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value=loaded) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = frozenset[str]().__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": "datadog"}) + + assert response.status_code == 200 + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"] + assert saved is loaded + assert saved.baseline == read_time + assert saved["litellm_settings"] == {"failure_callback": ["sentry"], "max_budget": 10} + + # --------------------------------------------------------------------------- # GET /get/config/callbacks # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 09d71cb4ead..f53e1931782 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -2090,3 +2090,36 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier() assert result is not None assert json.loads(result)["service_tier"] == "priority" + + +@pytest.mark.asyncio +async def test_async_data_generator_lets_other_tasks_run_during_a_stream_that_never_waits(monkeypatch): + """A stream whose chunks are already buffered never suspends on its own, so + the generator has to hand the loop back itself once it has held the loop + long enough: another task must get to run before the stream is over, and + not before the hold budget is used up.""" + _patch_logging_flags(monkeypatch) + order: list[str] = [] # mutable-ok: records the interleaving of the two tasks + clock: list[float] = [0.0] # mutable-ok: fake monotonic clock, advanced 20 ms per chunk + monkeypatch.setattr(ps.time, "monotonic", lambda: clock[0]) + + async def chunks(): + for index in range(6): + clock[0] += 0.02 + order.append(f"chunk{index}") + yield _simple_chunk(content=str(index)) + + async def other_task(): + order.append("other") + + other = asyncio.ensure_future(other_task()) + async for _ in async_data_generator( + response=chunks(), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ): + pass + await other + + assert "other" in order[:-1], order + assert order.index("other") >= 3, order diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 9e8910e44de..d5bacd9fd3c 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -437,15 +437,16 @@ def _decisions_provider_entry(provider: str) -> ProviderCreateInfo: return entry -def test_typesafe_provider_fields(): - typesafe: Final = _decisions_provider_entry("TypeSafe") +@pytest.mark.parametrize("display_name,provider", (("TypeSafe", LlmProviders.TYPESAFE), ("Nace", LlmProviders.NACE))) +def test_decisions_provider_fields(display_name: str, provider: LlmProviders) -> None: + entry: Final = _decisions_provider_entry(display_name) - assert typesafe.provider_display_name == "TypeSafe" - assert typesafe.litellm_provider == LlmProviders.TYPESAFE.value - assert typesafe.default_model_placeholder is not None - assert typesafe.default_model_placeholder.startswith("typesafe/") + assert entry.provider_display_name == display_name + assert entry.litellm_provider == provider.value + assert entry.default_model_placeholder is not None + assert entry.default_model_placeholder.startswith(f"{provider.value}/") - fields_by_key: Final = {f.key: f for f in typesafe.credential_fields} + fields_by_key: Final = {f.key: f for f in entry.credential_fields} assert fields_by_key["api_key"].required is True assert fields_by_key["api_key"].field_type == "password" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index b047c29243a..a7b955f5b84 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3454,10 +3454,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 = ( @@ -3474,6 +3475,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"}) @@ -3481,6 +3483,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?"}) @@ -3594,6 +3597,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/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index de2676bafaa..65dcce97b90 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4106,6 +4106,36 @@ def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] +@pytest.mark.parametrize("cached", [True, False]) +def test_get_logging_payload_gives_a_response_served_from_the_cache_a_unique_request_id(cached: bool): + """A parent call that returns an internal sub-call's cached response reuses the cached response's id, + and SpendLogs skips a duplicate request_id, so the row would be dropped unless its id is made unique.""" + from litellm.litellm_core_utils.hidden_params import set_hidden_params + + response_obj: Final = litellm.ModelResponse( + id="chatcmpl-cached-src", + model="gpt-5.6", + usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + set_hidden_params(response_obj, {"cache_hit": True, "response_cost": 0.0} if cached else {}) + now: Final = datetime.datetime.now(timezone.utc) + + payload: Final = get_logging_payload( + kwargs={ + "model": "openai/gpt-5.6", + "call_type": "aresponses", + "litellm_call_id": "parent-call-id", + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + }, + response_obj=response_obj, + start_time=now, + end_time=now, + ) + + assert payload["request_id"].startswith("chatcmpl-cached-src_cache_hit") is cached + assert (payload["request_id"] == "chatcmpl-cached-src") is not cached + + class TestSpendLogKeyRedaction: """Regression: plaintext API keys with Bearer prefix were stored in SpendLogs for failed requests (LIT-4121)""" diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py index 4e0889fd93d..adc4de8c8fa 100644 --- a/tests/unit/proxy/test__lazy_features.py +++ b/tests/unit/proxy/test__lazy_features.py @@ -401,3 +401,25 @@ def test_without_the_flag_the_warmup_route_registers_a_feature_and_returns_its_p assert body.paths["/alpha/list"]["get"].tags == ("alpha",) assert loaded_lazy_modules(app) == {features[0].module_path} assert "/alpha/list" in _paths(app) and "/beta/list" not in _paths(app) + + +def test_root_authorization_server_metadata_is_the_gateways_not_the_byok_connector_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Two features claim the root ``/.well-known/oauth-authorization-server``. The one a client + gets must be the gateway's own authorization server metadata, whose ``/token`` serves the + gateway's grants, not the BYOK connector form flow whose token endpoint is + ``/v1/mcp/oauth/token``; the connector keeps its path-appended document.""" + monkeypatch.setenv(FLAG, "false") + app: Final = FastAPI() + attach_lazy_features(app) + + with TestClient(app, base_url="https://gateway.example.com") as client: + root: Final = client.get("/.well-known/oauth-authorization-server") + connector: Final = client.get("/.well-known/oauth-authorization-server/v1/mcp/oauth") + + assert root.status_code == 200 + assert root.json()["issuer"] == "https://gateway.example.com" + assert root.json()["token_endpoint"] == "https://gateway.example.com/token" + assert connector.status_code == 200 + assert connector.json()["token_endpoint"] == "https://gateway.example.com/v1/mcp/oauth/token" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 1d197602543..c3556bbaeb7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1724,6 +1724,29 @@ class TestProxyBaseLLMRequestProcessing: assert "x-litellm-classifier-cost" not in headers + def test_get_custom_headers_drops_a_provider_header_http_cannot_carry(self): + user_api_key_dict: Final = MagicMock(spec=UserAPIKeyAuth) + user_api_key_dict.tpm_limit = None + user_api_key_dict.rpm_limit = None + user_api_key_dict.max_budget = None + user_api_key_dict.spend = 0 + + headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id="call-1", + **{ + "llm_provider-x-amzn-requestid": "req-1", + "llm_provider-x-provider-region": "東京リージョン", + "llm_provider-x-provider-city": "Zürich", + }, + ) + response: Final = Response() + response.headers.update(headers) + + assert response.headers["llm_provider-x-amzn-requestid"] == "req-1" + assert response.headers["llm_provider-x-provider-city"] == "Zürich" + assert "llm_provider-x-provider-region" not in response.headers + def test_get_cost_breakdown_from_logging_obj_helper(self): """ Test the helper function that extracts cost breakdown information. diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 7d70e058396..f7d1970fdae 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3,6 +3,7 @@ import json +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -153,6 +154,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-3-sonnet-20240229", "messages": [{"role": "user", "content": "Hello Claude!"}], @@ -172,7 +174,14 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -225,6 +234,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Hello GPT!"}], @@ -244,7 +254,14 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=True, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -1055,7 +1072,14 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1105,7 +1129,14 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1155,7 +1186,14 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1179,6 +1217,310 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +_GEMINI_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_bodies: tuple[bytes, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_bodies = (*self.request_bodies, request.content) + return httpx.Response(200, json={"totalTokens": 41}, request=request) + + +def _agent_turns_with_tool_results(rounds: int) -> list[dict[str, JsonValue]]: + return [ + *( + turn + for index in range(rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _GEMINI_TOOL_OUTPUT} + ], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ] + + +_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { + "name": "read_file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _gemini_deployment(model_name: str, litellm_params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **litellm_params}, + } + + +async def _count_through_the_messages_endpoint( + router: Router, request_options: Mapping[str, JsonValue] +) -> tuple[object, tuple[bytes, ...]]: + import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler + + peer: Final = _GeminiCountTokensPeer() + request_body: Final[dict[str, JsonValue]] = { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + **request_options, + } + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + + with ( + patch.object(proxy_server, "llm_router", router), + patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=request_body)), + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + ): + response: Final = await anthropic_count_tokens(mock_request, MagicMock()) + return response, peer.request_bodies + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ( + "requested_model", + "additional_drop_params", + "router_additional_drop_params", + "request_options", + "expected_tool_outputs_sent", + ), + ( + ("gemini-agent", None, None, {}, 1), + ("gemini-agent", ["context_management"], None, {}, 3), + ("gemini-agent", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini-agent", ["context_management"], None, {"additional_drop_params": []}, 1), + ("gemini-agent", None, ["context_management"], {}, 3), + ("gemini-agent", None, ["context_management"], {"additional_drop_params": []}, 1), + ("gemini/gemini-2.5-flash", None, None, {}, 1), + ("gemini/gemini-2.5-flash", ["context_management"], None, {}, 3), + ("gemini/gemini-2.5-flash", None, ["context_management"], {}, 3), + ("gemini/gemini-2.5-flash", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini/gemini-2.5-flash", ["context_management"], None, {"additional_drop_params": []}, 1), + ), + ids=( + "edits-applied", + "deployment-drops-context-management", + "request-drops-context-management", + "request-overrides-the-deployment-drop", + "router-default-drops-context-management", + "request-overrides-the-router-default-drop", + "litellm-model-name-edits-applied", + "litellm-model-name-deployment-drops-context-management", + "litellm-model-name-router-default-drops-context-management", + "litellm-model-name-request-drops-context-management", + "litellm-model-name-request-overrides-the-deployment-drop", + ), +) +async def test_gemini_messages_count_tokens_applies_the_context_edits_inference_applies( + requested_model: str, + additional_drop_params: list[str] | None, + router_additional_drop_params: list[str] | None, + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": additional_drop_params})], + default_litellm_params={"additional_drop_params": router_additional_drop_params}, + ) + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert len(sent_bodies) == 1 + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "model_list"), + ( + ( + "gemini-agent", + [ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", "order": 1}, + }, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "order": 2, + "additional_drop_params": ["context_management"], + }, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + {"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "test-key"}}, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ), + ids=("group-member-drops", "deployment-id-drops", "deployment-id-drops-beside-a-wildcard"), +) +async def test_gemini_messages_count_tokens_keeps_tool_results_when_a_routable_deployment_drops_context_management( + requested_model: str, model_list: list[dict[str, JsonValue]] +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=model_list), {"model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "make_router", "request_options", "expected_tool_outputs_sent"), + ( + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {}, + 1, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + default_fallbacks=["gemini-agent"], + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {})], + default_fallbacks=["gemini-agent"], + default_litellm_params={"additional_drop_params": ["context_management"]}, + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {"additional_drop_params": ["context_management"]}, + 3, + ), + ( + "flash-alias", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + model_group_alias={"flash-alias": "gemini/gemini-2.5-flash"}, + ), + {}, + 3, + ), + ), + ids=( + "default-fallback-edits-applied", + "default-fallback-deployment-drops-context-management", + "default-fallback-router-default-drops-context-management", + "default-fallback-request-drops-context-management", + "alias-to-a-litellm-model-deployment-drops-context-management", + ), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_deployment_the_router_falls_back_to( + requested_model: str, + make_router: Callable[[], Router], + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + make_router(), {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "completion_model", "expected_tool_outputs_sent"), + (("gemini-agent", "gemini-dropping", 3), ("gemini-dropping", "gemini-agent", 1)), + ids=("completion-model-drops-context-management", "completion-model-keeps-context-management"), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_completion_model_inference_sends_to( + requested_model: str, completion_model: str, expected_tool_outputs_sent: int +) -> None: + import litellm.proxy.proxy_server as proxy_server + + router: Final = Router( + model_list=[ + _gemini_deployment("gemini-agent", {}), + _gemini_deployment("gemini-dropping", {"additional_drop_params": ["context_management"]}), + ] + ) + with patch.object(proxy_server, "general_settings", {"completion_model": completion_model}): + response, sent_bodies = await _count_through_the_messages_endpoint(router, {"model": requested_model}) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + class _GeminiPart(BaseModel): model_config = ConfigDict(frozen=True) text: str = "" @@ -1224,6 +1566,48 @@ def _gemini_router(drop_params: bool = False) -> Router: ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: dict[str, JsonValue], + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], + litellm_drop_params: bool, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", deployment_params)], default_litellm_params=router_defaults + ) + with patch.object(litellm, "drop_params", litellm_drop_params): + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@pytest.mark.asyncio +async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None: + with patch.object(litellm, "drop_params", False): + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=[_gemini_deployment("gemini-agent", {})]), + {"model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT}, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert sent_bodies == () + + @pytest.mark.asyncio @pytest.mark.parametrize( "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] 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/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 6705e5ac8a5..da799a307a6 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3771,6 +3771,24 @@ class TestPtuCostAttributionUISetting: assert not mock_prisma.db.litellm_uisettings.upsert.called +class TestEnableChatUISetting: + def test_chat_ui_is_on_when_nothing_is_stored(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is True + + def test_admin_can_still_turn_chat_ui_off(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch, stored={"enable_chat_ui": False}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is False + + class TestApplyUserBudgetToTeamKeysUISetting: """``apply_user_budget_to_team_keys`` mirrors general_settings on every GET. 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/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 61b23428d62..025f1315aab 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -50,6 +50,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): litellm_completion_request: dict[str, object], instructions: str | None = None, request_kwargs: Mapping[str, object] | None = None, + previous_response_client_ids: tuple[str, ...] = (), ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index bb374b90f4e..d731bf5f41c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -13,7 +13,8 @@ kwargs lack the flag and these tests fail. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import SimpleNamespace from typing import Final from unittest.mock import patch @@ -25,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD @@ -223,3 +225,88 @@ async def test_bridged_follow_up_turn_keeps_the_addressed_response_id_off_the_pr ) assert isinstance(response, ResponsesAPIResponse) assert [item.type for item in response.output] == ["message"] + + +@pytest.mark.asyncio +async def test_bridged_turn_with_a_client_field_named_previous_response_client_ids_reaches_the_provider(): + provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + + response = await litellm.aresponses( + model="azure_ai/claude-sonnet-4-6", + api_base="https://fake-foundry-resource.services.ai.azure.com", + api_key="fake-api-key", + input="Double it", + client=client, + previous_response_client_ids=["resp_from_the_client_body"], + ) + + assert provider.request_body is not None, "the bridged turn never reached the provider" + assert isinstance(response, ResponsesAPIResponse) + + +class _StoredTurns: + def __init__(self, rows: tuple[Mapping[str, object], ...]) -> None: + self.rows: Final = rows + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict[str, object]]: + sessions: Final = {row["session_id"] for row in self.rows if row["request_id"] in request_ids} + return [dict(row) for row in self.rows if row["session_id"] in sessions] + + +def _stored_mcp_turn(request_id: str) -> Mapping[str, object]: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": "turn-one-session", + "proxy_server_request": {"input": "My favorite number is 7.", "model": "claude-sonnet-4-6"}, + "response": { + "id": request_id, + "object": "response", + "output": [ + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Noted.", "annotations": []}], + }, + {"type": "mcp_tools_fetched", "tools": [{"name": "demo-echo"}]}, + ], + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_under", ["litellm_encoded_id", "proxy_issued_id"]) +async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_holds(stored_under: str): + litellm_encoded_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure_ai", model_id="deployment-a", response_id="chatcmpl-turn-one" + ) + proxy_issued_id: Final = "resp_proxy-issued-turn-one" + stored_turns: Final = _StoredTurns( + (_stored_mcp_turn(litellm_encoded_id if stored_under == "litellm_encoded_id" else proxy_issued_id),) + ) + provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider)) + + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=stored_turns)): + await litellm.aresponses( + model="azure_ai/claude-sonnet-4-6", + api_base="https://fake-foundry-resource.services.ai.azure.com", + api_key="fake-api-key", + input="Double it", + previous_response_id=litellm_encoded_id, + client=client, + **{ADDRESSED_RESPONSE_ID_FIELD: proxy_issued_id}, + ) + + assert provider.request_body is not None, "the bridged turn never reached the provider" + assert [ + (message["role"], json.dumps(message["content"])) for message in provider.request_body["messages"] + ] == [ + ("user", json.dumps([{"type": "text", "text": "My favorite number is 7."}])), + ("assistant", json.dumps([{"type": "text", "text": "Noted."}])), + ("user", json.dumps([{"type": "text", "text": "Double it"}])), + ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 5db563be4b0..7a146c1b665 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,5 +1,7 @@ import json import sqlite3 +from collections.abc import Sequence +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -16,7 +18,9 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, _normalize_redacted_tool_call_arguments, ) +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import Message @@ -127,7 +131,7 @@ async def test_get_chat_completion_message_history_for_previous_response_id(): ) # Verify the mock was called with correct parameters - mock_get_spend_logs.assert_called_once_with(previous_response_id) + mock_get_spend_logs.assert_called_once_with(previous_response_id, ()) # Verify the returned ChatCompletionSession structure assert "messages" in result @@ -461,8 +465,8 @@ class _FakePrismaDB: self._results = list(results) self.calls = [] - async def query_raw(self, query, *args): - self.calls.append(args) + async def query_raw(self, query, request_ids): + self.calls.append(tuple(request_ids)) if not self._results: return [] return list(self._results.pop(0)) @@ -555,7 +559,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus ("assistant", "OK"), ] assert result["litellm_session_id"] == session_id - assert fake_prisma_client.db.calls == [(request_id,), (request_id,)] + assert [call[0] for call in fake_prisma_client.db.calls] == [request_id, request_id] @pytest.mark.asyncio @@ -623,7 +627,91 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): encoded_response_id ) - assert fake_prisma_client.db.calls == [(request_id,)] + assert len(fake_prisma_client.db.calls) == 1 + assert request_id in fake_prisma_client.db.calls[0] + + +class _StoredSpendLogs: + def __init__(self, rows: tuple[dict, ...]) -> None: + self.rows = rows + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict]: + sessions: Final = {row["session_id"] for row in self.rows if row["request_id"] in request_ids} + return [row for row in self.rows if row["session_id"] in sessions] + + +@pytest.mark.asyncio +async def test_message_history_replays_a_bridged_mcp_turn_stored_under_the_client_response_id(): + client_response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="anthropic", + model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", + response_id="chatcmpl-mcp-final-round", + ) + session_id: Final = "mcp-session" + stored: Final = _StoredSpendLogs( + ( + { + "request_id": "list-tools-row", + "call_type": "list_mcp_tools", + "session_id": session_id, + "proxy_server_request": {}, + "response": [{"name": "demo-echo"}], + }, + { + "request_id": client_response_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": {"input": "Call echo with PING.", "model": "claude-bridge"}, + "response": { + "id": client_response_id, + "object": "response", + "output": [ + { + "type": "reasoning", + "id": "rs_mcp-final-round", + "summary": [{"type": "summary_text", "text": "The echo tool returned PING."}], + }, + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "PING", "annotations": []}], + }, + *LiteLLM_Proxy_MCP_Handler.add_mcp_output_elements_to_response( + ResponsesAPIResponse(id=client_response_id, created_at=0, output=[]), + mcp_tools_fetched=[{"name": "demo-echo"}], + tool_results=[{"tool_name": "demo-echo", "result": "PING"}], + ).output, + ], + }, + }, + { + "request_id": "call-tool-row", + "call_type": "call_mcp_tool", + "session_id": session_id, + "proxy_server_request": {"name": "echo", "arguments": {"value": "PING"}}, + "response": {"content": [{"type": "text", "text": "PING"}], "is_error": False}, + }, + ) + ) + + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=stored)): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + client_response_id + ) + + assert [(message.get("role"), _text_of(message.get("content"))) for message in result["messages"]] == [ + ("user", "Call echo with PING."), + ("assistant", "PING"), + ] + assert result["messages"][-1].get("reasoning_content") == "The echo tool returned PING." + assert result["litellm_session_id"] == session_id + + +def _text_of(content: object) -> object: + if isinstance(content, list): + return "".join(part.get("text", "") for part in content if isinstance(part, dict)) + return content @pytest.mark.asyncio @@ -644,7 +732,7 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( ) assert spend_logs == [] - assert fake_prisma_client.db.calls == [("chatcmpl-does-not-exist",)] + assert len(fake_prisma_client.db.calls) == 1 def test_normalize_redacted_arguments_skips_custom_tool_calls(): @@ -801,7 +889,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)] + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) + return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] @pytest.mark.asyncio diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 4117c303940..1a3430224d1 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -11,6 +11,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve """ import json +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -978,6 +979,44 @@ def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelR ) +def _with_reasoning(chunk: ModelResponseStream, reasoning: str) -> ModelResponseStream: + choice: Final = chunk.choices[0] + delta: Final = Delta(**{**choice.delta.model_dump(), "reasoning_content": reasoning}) + return chunk.model_copy( + update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason=choice.finish_reason)]} + ) + + +def _reasoning_then_text_chunks() -> tuple[ModelResponseStream, ...]: + return (_reasoning_chunk("let me think"), _chunk("Hello"), _chunk("!", finish_reason="stop")) + + +def _reasoning_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: + return (_with_reasoning(_chunk("Hello!", finish_reason="stop"), "let me think"),) + + +def _signed_thinking_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: + chunk: Final = _chunk("Hello!", finish_reason="stop") + signed_block: Final = {"type": "thinking", "thinking": "", "signature": "sig_only"} + delta: Final = Delta(**{**chunk.choices[0].delta.model_dump(), "thinking_blocks": [signed_block]}) + return ( + chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}), + ) + + +StreamChunks = Callable[[], tuple[ModelResponseStream, ...]] + +REASONING_THEN_TEXT_STREAMS: Final = pytest.mark.parametrize( + "stream_chunks", + [_reasoning_then_text_chunks, _reasoning_and_text_in_one_chunk], + ids=["reasoning_then_text_chunks", "reasoning_and_text_in_one_chunk"], +) + + +def _event_text(events: list[BaseLiteLLMOpenAIResponseObject], event_type: ResponsesAPIStreamEvents) -> str: + return "".join(event.delta for event in events if getattr(event, "type", None) == event_type) + + def _annotation_only_chunk() -> ModelResponseStream: citation: Final = {"start_index": 0, "end_index": 2, "url": "https://example.com", "title": "Example"} return ModelResponseStream( @@ -1134,16 +1173,13 @@ async def test_signature_only_thinking_streams_a_replayable_reasoning_item(sync_ assert json.loads(reasoning_items[0].encrypted_content)[0]["signature"] == "sig_only" +@REASONING_THEN_TEXT_STREAMS @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool): - iterator: Final = _build_iterator( - [ - _reasoning_chunk("let me think"), - _chunk("Hello"), - _chunk("!", finish_reason="stop"), - ] - ) +async def test_reasoning_then_text_announces_message_item_before_text_events( + stream_chunks: StreamChunks, sync_mode: bool +): + iterator: Final = _build_iterator(stream_chunks()) events: Final = await _collect_events(iterator, sync_mode) @@ -1171,23 +1207,22 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): assert event.item.id in announced_message_ids assert saw_text_delta - assert "".join( - event.delta for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - ) == "Hello!" + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert _event_text(events, ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA) == "let me think" assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] +@pytest.mark.parametrize( + "stream_chunks", + [_reasoning_then_text_chunks, _reasoning_and_text_in_one_chunk, _signed_thinking_and_text_in_one_chunk], + ids=["reasoning_then_text_chunks", "reasoning_and_text_in_one_chunk", "signed_thinking_and_text_in_one_chunk"], +) +@pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_item_closes_before_message_item_opens(): - iterator: Final = _build_iterator( - [ - _reasoning_chunk("let me think"), - _chunk("Hello"), - _chunk("!", finish_reason="stop"), - ] - ) +async def test_reasoning_item_closes_before_message_item_opens(stream_chunks: StreamChunks, sync_mode: bool): + iterator: Final = _build_iterator(stream_chunks()) - events: Final = await _collect_events(iterator, sync_mode=False) + events: Final = await _collect_events(iterator, sync_mode) item_lifecycle: Final = [ (event.type, event.item.type) @@ -1201,6 +1236,7 @@ async def test_reasoning_item_closes_before_message_item_opens(): (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "message"), (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "message"), ] + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" @pytest.mark.parametrize("sync_mode", [True, False]) @@ -1231,6 +1267,27 @@ async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index( assert len(output_indexes_by_item_id) == len(set(output_indexes_by_item_id.values())) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_opens(sync_mode: bool): + iterator: Final = _build_iterator([_with_reasoning(_tool_call_chunk(finish_reason="tool_calls"), "let me think")]) + + events: Final = await _collect_events(iterator, sync_mode) + + item_lifecycle: Final = [ + (event.type, event.item.type) + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ] + assert item_lifecycle == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "reasoning"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "function_call"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "function_call"), + ] + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool): diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 745a48d0f6e..b9cf201a47a 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,18 +1,29 @@ +import asyncio import json import logging -from collections.abc import Mapping, Sequence -from types import MappingProxyType +import sys +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager, nullcontext +from datetime import datetime +from types import MappingProxyType, SimpleNamespace +from typing import Any, Final, List, cast +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from mcp.types import CallToolResult, ListToolsResult, TextContent from mcp.types import Tool as MCPTool -from typing import Final, List, Any, cast -from unittest.mock import AsyncMock, patch # Import required modules import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy._experimental.mcp_server import mcp_server_manager +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk +from litellm.proxy.utils import _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.types.llms.openai import ( ResponsesAPIResponse, @@ -20,6 +31,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, ToolParam, ) +from litellm.types.utils import StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1284,6 +1296,67 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): } +_ROUND_ONE_BODY: Final = MappingProxyType( + { + "id": "resp_round_one", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": ( + MappingProxyType( + { + "type": "function_call", + "id": "fc_round_one", + "call_id": "call_echo", + "name": "echo", + "arguments": '{"value": "PING"}', + "status": "completed", + } + ), + ), + "usage": MappingProxyType( + { + "input_tokens": 71, + "input_tokens_details": MappingProxyType({"cached_tokens": 32}), + "output_tokens": 19, + "output_tokens_details": MappingProxyType({"reasoning_tokens": 0}), + "total_tokens": 90, + } + ), + } +) +_ROUND_TWO_BODY: Final = MappingProxyType( + { + "id": "resp_round_two", + "object": "response", + "created_at": 2, + "status": "completed", + "model": "gpt-5.6", + "output": ( + MappingProxyType( + { + "type": "message", + "id": "msg_round_two", + "status": "completed", + "role": "assistant", + "content": (MappingProxyType({"type": "output_text", "text": "PING", "annotations": ()}),), + } + ), + ), + "usage": MappingProxyType( + { + "input_tokens": 149, + "input_tokens_details": MappingProxyType({"cached_tokens": 64}), + "output_tokens": 5, + "output_tokens_details": MappingProxyType({"reasoning_tokens": 0}), + "total_tokens": 154, + } + ), + } +) + + WEATHER_TOOL: Final = MCPTool( name="get_weather", description="Get weather for a city", @@ -1398,6 +1471,769 @@ SCRIPTED_PROVIDERS: Final = MappingProxyType( ) +def _round_cost(body: Mapping[str, object]) -> float: + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate(body), + model="gpt-5.6", + custom_llm_provider="openai", + call_type="aresponses", + ) + + +class _RoundRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[StandardLoggingPayload, ...] = () + self.failed: tuple[StandardLoggingPayload, ...] = () + self.hooked: tuple[object, ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: + self.hooked = (*self.hooked, result) + return kwargs, result + + +def _no_op_sync_callback( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + return None + + +async def _settle(condition: Callable[[], bool]) -> None: + for _ in range(200): + if condition(): + break + await asyncio.sleep(0.01) + await asyncio.sleep(0.2) + + +_ECHO_MCP_TOOL: Final = MappingProxyType( + {"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"} +) + + +def _serve_responses( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + + +def _serve_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, bodies: Sequence[Mapping[str, object]] +) -> respx.Route: + return _serve_responses( + respx_mock, monkeypatch, tuple(httpx.Response(200, content=json.dumps(body, default=dict)) for body in bodies) + ) + + +@contextmanager +def _echo_tool_served(tool_result: str) -> Iterator[None]: + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object( + LiteLLM_Proxy_MCP_Handler, + "execute_tool_calls", + AsyncMock(return_value=[{"tool_call_id": "call_echo", "result": tool_result, "name": "echo"}]), + ), + ): + yield + + +async def _auto_execute( + logging_obj: Logging, mcp_tool: Mapping[str, object], tool_result: str +) -> ResponsesAPIResponse: + with _echo_tool_served(tool_result): + response: Final = await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(mcp_tool)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_logging_obj=logging_obj, + ) + assert isinstance(response, ResponsesAPIResponse) + return response + + +async def _auto_execute_echo( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, logging_obj: Logging +) -> ResponsesAPIResponse: + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + return await _auto_execute(logging_obj, _ECHO_MCP_TOOL, "PING") + + +async def _auto_execute_through_mcp_server( + monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock, logging_obj: Logging +) -> None: + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + SimpleNamespace( + catalog=SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=call_tool, + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ), + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ): + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id=logging_obj.litellm_call_id, + litellm_logging_obj=logging_obj, + ) + + +def _rounds_logging_obj(recorder: _RoundRecorder, call_id: str) -> Logging: + return Logging( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + stream=False, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_success_callbacks=[recorder, _no_op_sync_callback], + dynamic_async_failure_callbacks=[recorder], + ) + + +def _output_types(output: Sequence[object]) -> tuple[object, ...]: + return tuple(item.get("type") if isinstance(item, Mapping) else getattr(item, "type", None) for item in output) + + +def _assert_one_row_bills_both_rounds(logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse) -> None: + assert len(logged) == 1 + row: Final = logged[0] + both_rounds: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) + assert row["id"] == client.id + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["response_cost"] == pytest.approx(both_rounds) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(both_rounds) + assert client.hidden_params["response_cost"] == pytest.approx(both_rounds) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_logs_the_client_response_billed_for_both_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + response: Final = await _auto_execute_echo( + respx_mock, monkeypatch, _rounds_logging_obj(recorder, "mcp-rounds-immediate") + ) + + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_both_rounds(recorder.logged, response) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-rounds-deferred") + logging_obj.defer_async_logging = True + response: Final = await _auto_execute_echo(respx_mock, monkeypatch, logging_obj) + + enqueue_deferred: Final = logging_obj.enqueue_deferred_logging + assert enqueue_deferred is not None + enqueue_deferred() + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_both_rounds(recorder.logged, response) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_hands_sync_callbacks_the_client_response( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _RoundRecorder() + response: Final = await _auto_execute_echo(respx_mock, monkeypatch, _rounds_logging_obj(recorder, "mcp-rounds-sync")) + + await _settle(lambda: len(recorder.hooked) > 0) + + assert len(recorder.hooked) == 1 + hooked: Final = recorder.hooked[0] + assert isinstance(hooked, ResponsesAPIResponse) + assert hooked.id == response.id + assert _output_types(hooked.output) == _output_types(response.output) + assert hooked.usage is not None + assert (hooked.usage.model_dump().get("prompt_tokens"), hooked.usage.model_dump().get("completion_tokens")) == ( + 71 + 149, + 19 + 5, + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_follow_up_fails( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, + monkeypatch, + ( + httpx.Response(200, content=json.dumps(_ROUND_ONE_BODY, default=dict)), + httpx.Response(400, json=context_window_error), + ), + ) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + recorder: Final = _RoundRecorder() + + with pytest.raises(litellm.BadRequestError): + await _auto_execute_through_mcp_server( + monkeypatch, call_tool, _rounds_logging_obj(recorder, "mcp-rounds-follow-up-fails") + ) + await _settle(lambda: len(recorder.failed) > 0) + + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + assert recorder.failed[0]["response_cost"] > 0 + assert call_tool.await_count == 1 + assert call_tool.await_args is not None + tool_logging_obj: Final = call_tool.await_args.kwargs["litellm_logging_obj"] + assert tool_logging_obj.litellm_call_id == f"{recorder.failed[0]['id']}_call_echo" + assert tool_logging_obj.litellm_trace_id == recorder.failed[0]["trace_id"] + + +_RETRIED_ROUND_ONE: Final = MappingProxyType( + { + **_ROUND_ONE_BODY, + "id": "resp_retried_round_one", + "usage": { + "input_tokens": 500, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 40, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 540, + }, + } +) +_SERVER_ERROR: Final = MappingProxyType( + {"error": {"message": "upstream overloaded", "type": "server_error", "code": "server_error"}} +) + + +def _ok(body: Mapping[str, object]) -> httpx.Response: + return httpx.Response(200, content=json.dumps(body, default=dict)) + + +def _server_error() -> httpx.Response: + return httpx.Response(500, content=json.dumps(_SERVER_ERROR, default=dict)) + + +def _router_retrying_once() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-mcp-rounds", + "api_base": "https://mcp-rounds.example.com/v1", + }, + } + ], + num_retries=1, + retry_after=0, + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_retried_after_a_failed_follow_up_bills_each_attempt_its_own_rounds( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses( + respx_mock, + monkeypatch, + (_ok(_ROUND_ONE_BODY), _server_error(), _ok(_RETRIED_ROUND_ONE), _ok(_ROUND_TWO_BODY)), + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + response: Final = await _router_retrying_once().aresponses( + model="gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-rounds-retried"), + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert isinstance(response, ResponsesAPIResponse) + assert response.usage is not None + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (71, 19) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert row["id"] == response.id + assert (row["prompt_tokens"], row["completion_tokens"]) == (500 + 149, 40 + 5) + assert (row["prompt_tokens"], row["completion_tokens"]) == (response.usage.input_tokens, response.usage.output_tokens) + assert row["response_cost"] == pytest.approx(_round_cost(_RETRIED_ROUND_ONE) + _round_cost(_ROUND_TWO_BODY)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retried_first_rounds", "billed_tokens"), + [ + pytest.param((), (71, 19), id="the retry fails its first call"), + pytest.param((_RETRIED_ROUND_ONE,), (71 + 500, 19 + 40), id="the retry fails its follow-up"), + ], +) +async def test_non_stream_auto_execute_failing_every_retry_bills_each_paid_first_round_once_on_the_request_failure( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + retried_first_rounds: tuple[Mapping[str, object], ...], + billed_tokens: tuple[int, int], +): + _serve_responses( + respx_mock, + monkeypatch, + (_ok(_ROUND_ONE_BODY), _server_error(), *(_ok(body) for body in retried_first_rounds), _server_error()), + ) + logging_obj: Final = _rounds_logging_obj(_RoundRecorder(), "mcp-rounds-every-retry-fails") + + with _echo_tool_served("PING"), pytest.raises(litellm.InternalServerError): + await _router_retrying_once().aresponses( + model="gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + litellm_logging_obj=logging_obj, + ) + + failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + billed_usage: Final = failure_row_fields["combined_usage_object"] + assert isinstance(billed_usage, Usage) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == billed_tokens + assert failure_row_fields["response_cost"] == pytest.approx( + _round_cost(_ROUND_ONE_BODY) + sum(_round_cost(body) for body in retried_first_rounds) + ) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_tool_log_continues_the_request_trace_without_renaming_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + recorder: Final = _RoundRecorder() + + await _auto_execute_through_mcp_server(monkeypatch, call_tool, _rounds_logging_obj(recorder, "mcp-rounds-trace")) + await _settle(lambda: len(recorder.logged) > 0) + + assert call_tool.await_args is not None + tool_metadata: Final = call_tool.await_args.kwargs["litellm_logging_obj"].litellm_params["metadata"] + assert tool_metadata["existing_trace_id"] == recorder.logged[0]["trace_id"] + + +class _ListingRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.listings: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs.get("call_type") == "list_mcp_tools": + self.listings = (*self.listings, kwargs) + + +def _serve_the_demo_mcp_listing(monkeypatch: pytest.MonkeyPatch) -> _ListingRecorder: + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + demo_server: Final = SimpleNamespace( + server_id="demo", name="demo", alias="demo", server_name="demo", auth_type=None, extra_headers=None + ) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + SimpleNamespace( + catalog=SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + get_allowed_mcp_servers=AsyncMock(return_value=["demo"]), + get_mcp_servers_from_ids=MagicMock(return_value=[demo_server]), + ), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", + AsyncMock(return_value=[demo_server]), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.catalog.get_filtered_server_tools", + AsyncMock(return_value=(ListToolsResult(tools=[echo_tool]), ServerListOk(tool_count=1))), + ) + listing_recorder: Final = _ListingRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [listing_recorder]) + return listing_recorder + + +async def _assert_the_listing_logs_in_the_request_trace( + listing_recorder: _ListingRecorder, request_recorder: _RoundRecorder +) -> None: + await _settle(lambda: len(request_recorder.logged) > 0 and len(listing_recorder.listings) > 0) + request_trace_id: Final = request_recorder.logged[0]["trace_id"] + assert len(listing_recorder.listings) == 1 + listing: Final = listing_recorder.listings[0] + listing_row: Final = cast(StandardLoggingPayload, listing["standard_logging_object"]) + assert listing_row["trace_id"] == request_trace_id + assert cast(Mapping[str, Mapping[str, object]], listing["litellm_params"])["metadata"][ + "existing_trace_id" + ] == request_trace_id + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_tool_listing_log_continues_the_request_trace( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + request_recorder: Final = _RoundRecorder() + + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id="mcp-listing-trace", + litellm_trace_id="router-trace", + litellm_logging_obj=_rounds_logging_obj(request_recorder, "mcp-listing-trace"), + ) + + await _assert_the_listing_logs_in_the_request_trace(listing_recorder, request_recorder) + + +@pytest.mark.asyncio +async def test_bridged_auto_execute_tool_listing_log_continues_the_request_trace( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://mcp-rounds.example.com/v1/messages").mock( + side_effect=[ + httpx.Response( + 200, + json=_anthropic_message( + "msg_tool_round", + ({"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}},), + "tool_use", + ), + ), + httpx.Response(200, json=_anthropic_message("msg_answer", ({"type": "text", "text": "PING"},), "end_turn")), + ] + ) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + request_recorder: Final = _RoundRecorder() + + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com", + litellm_call_id="mcp-bridged-listing-trace", + litellm_trace_id="router-trace", + litellm_logging_obj=_rounds_logging_obj(request_recorder, "mcp-bridged-listing-trace"), + ) + + await _assert_the_listing_logs_in_the_request_trace(listing_recorder, request_recorder) + + +class _CallRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[Mapping[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.logged = (*self.logged, kwargs) + + def of(self, call_type: str) -> tuple[Mapping[str, object], ...]: + return tuple(kwargs for kwargs in self.logged if kwargs.get("call_type") == call_type) + + +def _langfuse_trace_of(logged: Mapping[str, object]) -> object: + litellm_params: Final = cast(Mapping[str, object], logged["litellm_params"]) + return langfuse_trace_id( + metadata=langfuse_trace_metadata(litellm_params["metadata"], litellm_params.get("proxy_server_request")), + standard_trace_id=cast(StandardLoggingPayload, logged["standard_logging_object"])["trace_id"], + litellm_trace_id=logged.get("litellm_trace_id"), + litellm_call_id=logged.get("litellm_call_id"), + proxy_server_request=litellm_params.get("proxy_server_request"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body_metadata", + [None, {"trace_id": "client-trace", "session_id": "client-session"}], + ids=["no_body_metadata", "body_metadata_with_its_own_trace_and_session"], +) +async def test_auto_execute_mcp_logs_join_the_langfuse_trace_of_a_session_header_request( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body_metadata: dict[str, str] | None +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + listing_recorder: Final = _serve_the_demo_mcp_listing(monkeypatch) + call_recorder: Final = _CallRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [listing_recorder, call_recorder]) + session_id: Final = "sess-abc12345" + + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + litellm_call_id="mcp-session-trace", + litellm_trace_id=session_id, + litellm_session_id=session_id, + metadata=body_metadata, + litellm_metadata={"trace_id": session_id, "session_id": session_id}, + proxy_server_request={"headers": {"x-litellm-session-id": session_id}}, + ) + await _settle(lambda: len(call_recorder.of("aresponses")) > 0 and len(listing_recorder.listings) > 0) + + parent_trace_id: Final = _langfuse_trace_of(call_recorder.of("aresponses")[0]) + call_tool: Final = cast(AsyncMock, mcp_server_manager.global_mcp_server_manager.call_tool) + assert call_tool.await_args is not None + tool_logging_obj: Final = cast(Logging, call_tool.await_args.kwargs["litellm_logging_obj"]) + mcp_log_traces: Final = ( + cast(Mapping[str, Mapping[str, object]], listing_recorder.listings[0]["litellm_params"])["metadata"][ + "existing_trace_id" + ], + cast(Mapping[str, Mapping[str, object]], tool_logging_obj.model_call_details["litellm_params"])["metadata"][ + "existing_trace_id" + ], + ) + assert parent_trace_id == "mcp-session-trace" + assert mcp_log_traces == (parent_trace_id, parent_trace_id) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_the_paid_first_round_when_the_client_disconnects_during_a_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY,)) + logging_obj: Final = _rounds_logging_obj(_RoundRecorder(), "mcp-rounds-disconnect") + + with pytest.raises(asyncio.CancelledError): + await _auto_execute_through_mcp_server( + monkeypatch, AsyncMock(side_effect=asyncio.CancelledError()), logging_obj + ) + + failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + billed_usage: Final = failure_row_fields["combined_usage_object"] + assert isinstance(billed_usage, Usage) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == (71, 19) + assert failure_row_fields["response_cost"] == pytest.approx(_round_cost(_ROUND_ONE_BODY)) + + +@pytest.mark.asyncio +async def test_non_stream_auto_execute_bills_nothing_for_rounds_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + """A cached final round hands back the id of the turn that cached it, so that turn's row is logged under a + unique id prefixed by the client's, and as a cache hit only when no round was paid.""" + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY, _ROUND_TWO_BODY)) + warm: Final = _RoundRecorder() + await _auto_execute(_rounds_logging_obj(warm, "mcp-cache-warm"), _ECHO_MCP_TOOL, "PING") + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + first_cached_response: Final = await _auto_execute( + _rounds_logging_obj(first_cached, "mcp-cache-first-round"), + {**_ECHO_MCP_TOOL, "server_description": "echo server"}, + "PONG", + ) + all_cached: Final = _RoundRecorder() + all_cached_response: Final = await _auto_execute( + _rounds_logging_obj(all_cached, "mcp-cache-every-round"), + {**_ECHO_MCP_TOOL, "server_description": "echo server, same tools"}, + "PONG", + ) + await _settle(lambda: len(first_cached.logged) > 0 and len(all_cached.logged) > 0) + + assert provider.call_count == 3 + for logged, client in ((first_cached.logged, first_cached_response), (all_cached.logged, all_cached_response)): + assert len(logged) == 1 + assert logged[0]["id"].startswith(client.id) + assert _output_types(logged[0]["response"]["output"]) == _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (71 + 149, 19 + 5) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert first_cached.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert first_cached_response.hidden_params["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY)) + assert all_cached.logged[0]["response_cost"] == 0.0 + assert all_cached_response.hidden_params["response_cost"] == 0.0 + assert first_cached.logged[0]["id"] == first_cached_response.id + assert first_cached.logged[0]["cache_hit"] is None + assert all_cached.logged[0]["id"].startswith(f"{all_cached_response.id}_cache_hit") + assert all_cached.logged[0]["cache_hit"] is True + + +_PRIOR_TURN: Final = MappingProxyType( + { + "request_id": "chatcmpl-prior-turn", + "session_id": "prior-session", + "proxy_server_request": {"input": "My favorite color is teal."}, + "response": { + "id": "chatcmpl-prior-turn", + "object": "chat.completion", + "created": 1767225600, + "model": "claude-haiku-4-5", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Noted."}}], + }, + } +) + + +def _anthropic_message(message_id: str, content: Sequence[Mapping[str, object]], stop: str) -> dict[str, object]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [dict(block) for block in content], + "stop_reason": stop, + "usage": {"input_tokens": 50, "output_tokens": 10}, + } + + +class _SpendLogs: + def __init__(self, rows: Sequence[Mapping[str, object]]) -> None: + self.rows: Final = tuple(rows) + self.asked: tuple[frozenset[str], ...] = () + + async def query_raw(self, query: str, request_ids: Sequence[str]) -> list[dict[str, object]]: + self.asked = (*self.asked, frozenset(request_ids)) + return [dict(row) for row in self.rows if row["request_id"] in request_ids] + + +@pytest.mark.asyncio +async def test_bridged_auto_execute_follow_up_continues_the_client_conversation_in_its_session( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + spend_logs: Final = _SpendLogs((_PRIOR_TURN,)) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(proxy_logging_obj=None, prisma_client=SimpleNamespace(db=spend_logs), disable_spend_logs=False), + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + thinking: Final = {"type": "thinking", "thinking": "Echo PING back.", "signature": "sig-round-one"} + provider: Final = respx_mock.post("https://mcp-rounds.example.com/v1/messages").mock( + side_effect=[ + httpx.Response( + 200, + json=_anthropic_message( + "msg_tool_round", + (thinking, {"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}), + "tool_use", + ), + ), + httpx.Response(200, json=_anthropic_message("msg_answer", ({"type": "text", "text": "PING"},), "end_turn")), + ] + ) + echo_tool: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + execute: Final = AsyncMock(return_value=[{"tool_call_id": "toolu_echo", "result": "PING", "name": "echo"}]) + recorder: Final = _RoundRecorder() + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([echo_tool], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), + ): + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + previous_response_id="chatcmpl-prior-turn", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com", + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-rounds-bridged"), + ) + await _settle(lambda: len(recorder.logged) > 0) + + follow_up: Final = json.loads(provider.calls[1].request.content)["messages"] + assert spend_logs.asked == (frozenset({"chatcmpl-prior-turn"}),) * 2 + assert [message["role"] for message in follow_up] == ["user", "assistant", "user", "assistant", "user"] + assert json.dumps(follow_up).count("My favorite color is teal.") == 1 + assert [block["type"] for block in follow_up[3]["content"]] == ["thinking", "tool_use"] + assert follow_up[3]["content"][0]["signature"] == "sig-round-one" + assert execute.await_args is not None + assert execute.await_args.kwargs["litellm_trace_id"] == "prior-session" + assert recorder.logged[0]["trace_id"] == "prior-session" + + def _event_type(event: object) -> str: event_type: Final = getattr(event, "type", "") return str(getattr(event_type, "value", event_type)) diff --git a/tests/unit/responses/mcp/test_round_billing.py b/tests/unit/responses/mcp/test_round_billing.py new file mode 100644 index 00000000000..993b38de16a --- /dev/null +++ b/tests/unit/responses/mcp/test_round_billing.py @@ -0,0 +1,218 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest + +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_round_usage, summed_cost_breakdown +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import CostBreakdown + + +def _round(response_id: str, response_cost: object, served_from_cache: bool = False) -> ResponsesAPIResponse: + response: Final = ResponsesAPIResponse.model_validate( + MappingProxyType( + { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": (), + "usage": MappingProxyType({"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}), + } + ) + ) + response.hidden_params["response_cost"] = response_cost + if served_from_cache: + response.hidden_params["cache_hit"] = True + return response + + +def test_summed_cost_breakdown_adds_each_charge_and_keeps_the_final_round_rates(): + first: Final[CostBreakdown] = { + "input_cost": 0.001, + "output_cost": 0.002, + "total_cost": 0.003, + "additional_costs": {"web_search": 0.01}, + "discount_percent": 0.1, + "margin_percent": 0.2, + "service_tier": "default", + } + final: Final[CostBreakdown] = { + "input_cost": 0.004, + "output_cost": 0.005, + "total_cost": 0.009, + "additional_costs": {"web_search": 0.01, "file_search": 0.0025}, + "discount_percent": 0.1, + "margin_percent": 0.2, + "service_tier": "default", + } + + summed: Final = summed_cost_breakdown( + first=_round("resp_first", 0.003), first_breakdown=first, final=_round("resp_final", 0.009), final_breakdown=final + ) + + assert summed is not None + assert summed["input_cost"] == pytest.approx(0.005) + assert summed["output_cost"] == pytest.approx(0.007) + assert summed["total_cost"] == pytest.approx(0.012) + assert summed["additional_costs"] == {"web_search": pytest.approx(0.02), "file_search": pytest.approx(0.0025)} + assert (summed["discount_percent"], summed["margin_percent"]) == (0.1, 0.2) + assert summed["service_tier"] == "default" + + +def test_summed_cost_breakdown_is_none_when_the_final_round_was_not_repriced(): + breakdown: Final[CostBreakdown] = {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003} + first: Final = _round("resp_first", 0.003) + final: Final = _round("resp_final", 0.003) + + assert summed_cost_breakdown(first=first, first_breakdown=breakdown, final=final, final_breakdown=breakdown) is None + assert summed_cost_breakdown(first=first, first_breakdown=None, final=final, final_breakdown=breakdown) is None + assert ( + summed_cost_breakdown( + first=_round("resp_first", 0.003, served_from_cache=True), + first_breakdown=breakdown, + final=final, + final_breakdown=breakdown, + ) + is None + ) + + +def test_a_final_round_served_from_the_response_cache_keeps_the_paid_first_round_breakdown(): + breakdown: Final[CostBreakdown] = {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003} + + assert ( + summed_cost_breakdown( + first=_round("resp_first", 0.003), + first_breakdown=breakdown, + final=_round("resp_final", 0.0, served_from_cache=True), + final_breakdown=breakdown, + ) + == breakdown + ) + + +def test_merged_round_usage_sums_token_details_and_keeps_provider_flags(): + first: Final = ResponseAPIUsage.model_validate( + MappingProxyType( + { + "input_tokens": 71, + "input_tokens_details": MappingProxyType({"cached_tokens": 32}), + "output_tokens": 19, + "total_tokens": 90, + "is_byok": True, + } + ) + ) + final: Final = ResponseAPIUsage.model_validate( + MappingProxyType( + { + "input_tokens": 149, + "input_tokens_details": MappingProxyType({"cached_tokens": 64}), + "output_tokens": 5, + "total_tokens": 154, + "is_byok": True, + } + ) + ) + + merged: Final = merged_round_usage(first, final) + + assert merged is not None + assert (merged.input_tokens, merged.output_tokens, merged.total_tokens) == (220, 24, 244) + assert merged.input_tokens_details is not None + assert merged.input_tokens_details.cached_tokens == 96 + assert merged.model_dump()["is_byok"] is True + + +def test_billed_for_every_round_leaves_the_cost_to_the_caller_when_a_round_has_no_price(): + billed: Final = billed_for_every_round( + first=_round("resp_first", None), first_breakdown=None, final=_round("resp_final", 0.004), final_breakdown=None + ) + + assert billed.id == "resp_final" + assert billed.hidden_params["response_cost"] is None + + +def test_billed_for_every_round_does_not_rebill_the_final_round_object(): + final: Final = _round("resp_final", 0.004) + + billed: Final = billed_for_every_round( + first=_round("resp_first", 0.001), first_breakdown=None, final=final, final_breakdown=None + ) + + assert billed.hidden_params["response_cost"] == pytest.approx(0.005) + assert final.hidden_params["response_cost"] == 0.004 + + +@pytest.mark.parametrize( + ("final_cached", "request_cost", "request_margin"), + [ + (False, 0.012, 0.002), + (True, 0.0032, 0.0012), + ], +) +def test_the_fixed_margin_is_charged_once_per_request_across_charged_rounds( + final_cached: bool, request_cost: float, request_margin: float +): + first_breakdown: Final[CostBreakdown] = { + "input_cost": 0.0015, + "output_cost": 0.0005, + "total_cost": 0.0032, + "margin_percent": 0.1, + "margin_fixed_amount": 0.001, + "margin_total_amount": 0.0012, + } + final_breakdown: Final[CostBreakdown] = { + "input_cost": 0.006, + "output_cost": 0.002, + "total_cost": 0.0098, + "margin_percent": 0.1, + "margin_fixed_amount": 0.001, + "margin_total_amount": 0.0018, + } + first: Final = _round("resp_first", 0.0032) + final: Final = _round("resp_final", 0.0 if final_cached else 0.0098, served_from_cache=final_cached) + + billed: Final = billed_for_every_round( + first=first, first_breakdown=first_breakdown, final=final, final_breakdown=final_breakdown + ) + breakdown: Final = summed_cost_breakdown( + first=first, first_breakdown=first_breakdown, final=final, final_breakdown=final_breakdown + ) + + assert billed.hidden_params["response_cost"] == pytest.approx(request_cost) + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(request_cost) + assert breakdown["margin_fixed_amount"] == pytest.approx(0.001) + assert breakdown["margin_total_amount"] == pytest.approx(request_margin) + + +@pytest.mark.parametrize( + ("first_cached", "final_cached", "kept_breakdown"), + [ + (True, False, "final"), + (False, True, "first"), + (True, True, None), + ], +) +def test_a_round_served_from_the_response_cache_drops_out_of_the_breakdown( + first_cached: bool, final_cached: bool, kept_breakdown: str | None +): + breakdowns: Final[Mapping[str, CostBreakdown]] = MappingProxyType( + { + "first": {"input_cost": 0.001, "output_cost": 0.002, "total_cost": 0.003}, + "final": {"input_cost": 0.004, "output_cost": 0.005, "total_cost": 0.009}, + } + ) + + breakdown: Final = summed_cost_breakdown( + first=_round("resp_first", 0.003, served_from_cache=first_cached), + first_breakdown=breakdowns["first"], + final=_round("resp_final", 0.009, served_from_cache=final_cached), + final_breakdown=breakdowns["final"], + ) + + assert breakdown == (None if kept_breakdown is None else breakdowns[kept_breakdown]) diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 685ea46547e..380aa573b3c 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -14,6 +14,7 @@ import pytest from litellm import Router from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig import sys, os, asyncio, time, random @@ -232,7 +233,11 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( """ monkeypatch.setattr(litellm, "tag_budget_config", tag_budget_config) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) - limiter = RouterBudgetLimiting(dual_cache=DualCache(), clock=charged_at.timestamp, **limiter_config) + limiter = RouterBudgetLimiting( + dual_cache=DualCache(in_memory_cache=InMemoryCache(clock=charged_at.timestamp)), + clock=charged_at.timestamp, + **limiter_config, + ) await limiter.dual_cache.async_set_cache(key=start_time_key, value=opened_at.timestamp()) await _log_success( @@ -243,9 +248,8 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( }, ) - assert limiter.redis_increment_operation_queue == [ - {"key": spend_key, "increment_value": 0.25, "ttl": int((closes_at - charged_at).total_seconds())} - ] + assert await limiter.dual_cache.async_get_cache(spend_key) == 0.25 + assert limiter.dual_cache.in_memory_cache.ttl_dict[spend_key] == closes_at.timestamp() @pytest.mark.asyncio @@ -362,24 +366,18 @@ async def test_new_budget_window_sets_spend_and_start_time() -> None: @pytest.mark.asyncio -async def test_spend_increment_updates_memory_and_queues_redis_operation() -> None: - provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) - await provider_budget.dual_cache.async_set_cache(key="provider_spend:openai:1d", value=1.0, ttl=86400) - - await provider_budget._increment_spend_in_current_window( - spend_key="provider_spend:openai:1d", - response_cost=0.5, - ttl=86400, +async def test_redis_less_limiter_tracks_spend_without_queueing_redis_increments(disable_budget_sync) -> None: + """With no Redis there is nothing to flush to, so queued increments would be retained for the life of the process.""" + limiter: Final = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, ) - assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 1.5 - assert provider_budget.redis_increment_operation_queue == [ - { - "key": "provider_spend:openai:1d", - "increment_value": 0.5, - "ttl": 86400, - } - ] + for _ in range(3): + await _log_success(limiter, _success_kwargs(provider_in_litellm_params="openai", provider_in_payload="openai")) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(0.75) + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio 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..2fbe6ec294b 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -229,6 +229,62 @@ def test_response_cost_calculator_custom_pricing_ignores_upstream_response_cost_ assert cost == pytest.approx(expected_cost) +def _header_cost_chat_response(additional_headers: dict[str, str]) -> ModelResponse: + response: Final = ModelResponse( + model="cc-header-cost-chat", + choices=[Choices(message=Message(content="ok"))], + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + response._hidden_params = {"additional_headers": additional_headers} + return response + + +def _cost_with_upstream_headers(monkeypatch: pytest.MonkeyPatch, additional_headers: dict[str, str]) -> float: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.register_model( + model_cost={ + "cc-header-cost-chat": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + return response_cost_calculator( + response_object=_header_cost_chat_response(additional_headers), + model="cc-header-cost-chat", + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + +@pytest.mark.parametrize("header_cost", ("-5", "not-a-number", "nan", "inf")) +def test_response_cost_calculator_prices_from_the_map_when_upstream_cost_header_is_not_billable( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, +) -> None: + price_map_cost: Final = _cost_with_upstream_headers(monkeypatch, {}) + + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert price_map_cost > 0 + assert cost == pytest.approx(price_map_cost) + + +@pytest.mark.parametrize(("header_cost", "expected_cost"), (("0", 0.0), ("0.25", 0.25))) +def test_response_cost_calculator_bills_a_finite_non_negative_upstream_cost_header( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, + expected_cost: float, +) -> None: + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert cost == expected_cost + + def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): class MockResponse(BaseModel): pass @@ -6426,6 +6482,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_main.py b/tests/unit/test_main.py index 14cea8f0c2c..bb0e8db7f1c 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -21,15 +21,22 @@ import httpx import pytest import respx from openai import APITimeoutError, AsyncOpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion import litellm -from litellm import acompletion, completion -from litellm import acompletion_with_retries, aresponses_with_retries, completion_with_retries, responses_with_retries +from litellm import ( + acompletion, + acompletion_with_retries, + aresponses_with_retries, + completion, + completion_with_retries, + responses_with_retries, +) from litellm import main as litellm_main -from litellm.constants import CONTROL_OPTIONS_KEY from litellm.caching.base_cache import BaseCache from litellm.caching.caching import Cache +from litellm.constants import CONTROL_OPTIONS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -130,7 +137,7 @@ def test_completion_missing_role(openai_api_response): client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -142,7 +149,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) + client, "post", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", @@ -218,7 +225,7 @@ async def test_url_with_format_param_openai(model, sync_mode): } ], } - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: if sync_mode: response = completion(**args, client=client) @@ -232,7 +239,7 @@ async def test_url_with_format_param_openai(model, sync_mode): print(mock_client.call_args.kwargs) - json_str = json.dumps(mock_client.call_args.kwargs) + json_str = json.dumps(mock_client.call_args.kwargs["body"]) assert "format" not in json_str @@ -3005,9 +3012,10 @@ def test_stream_chunk_builder_thinking_blocks(): assert response.choices[0].message.thinking_blocks is not None -from litellm.llms.openai.openai import OpenAIChatCompletion import traceback +from litellm.llms.openai.openai import OpenAIChatCompletion + user_message = "Write a short poem about the sky" messages = [{"content": user_message, "role": "user"}] @@ -3777,7 +3785,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3790,7 +3798,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3805,7 +3813,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): client = AsyncOpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: await litellm.acompletion( model="openai/gpt-4o", @@ -3818,7 +3826,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3832,7 +3840,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3843,7 +3851,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert "store" not in request_body assert "prompt_cache_key" not in request_body @@ -4038,8 +4046,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): from openai import OpenAI client = OpenAI(api_key="fake-key") - raw_response = client.chat.completions.with_raw_response - with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + with patch.object(client, "post") as mock_create, contextlib.suppress(Exception): litellm.completion( model="gpt-4o", messages=[{"role": "user", "content": "hello"}], @@ -4050,7 +4057,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): def _openai_chat_create_kwargs(client, **completion_kwargs): - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: with contextlib.suppress(Exception): litellm.completion( messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], @@ -4075,7 +4082,7 @@ def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4087,7 +4094,7 @@ def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4100,7 +4107,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: with contextlib.suppress(Exception): await litellm.acompletion( model="gpt-5.6", @@ -4111,7 +4118,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ ) mock_create.assert_called_once() - request_body = mock_create.call_args.kwargs + request_body = mock_create.call_args.kwargs["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4123,12 +4130,12 @@ def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6() from openai import OpenAI client = OpenAI(api_key="fake-api-key") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") + request = _openai_chat_create_kwargs(client, model="gpt-5.6") - assert request_body["messages"][0]["content"] == [ + assert request["body"]["messages"][0]["content"] == [ {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} ] - assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "implicit"} + assert request["options"]["extra_json"]["prompt_cache_options"] == {"mode": "implicit"} _SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" @@ -5141,6 +5148,7 @@ def test_drop_params_false_still_rejects_an_invalid_stream_chunk_size() -> None: def test_acompletion_params(): import inspect + from litellm.types.completion import CompletionRequest acompletion_params_odict = inspect.signature(acompletion).parameters @@ -5171,7 +5179,7 @@ def test_acompletion_params(): def _openai_mock_response(*args: object, **kwargs: object) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -5206,7 +5214,7 @@ def test_null_role_response(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, "create", side_effect=_openai_mock_response + openai_client, "post", side_effect=_openai_mock_response ) as mock_response: response = litellm.completion( model="gpt-3.5-turbo", @@ -5296,7 +5304,7 @@ def test_completion_perplexity_api(): pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} new_response.parse.return_value = pydantic_obj @@ -5305,8 +5313,7 @@ def test_completion_perplexity_api(): openai_client = OpenAI() with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: # litellm.set_verbose= True @@ -5350,7 +5357,7 @@ async def test_openai_compatible_custom_api_base(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -5368,7 +5375,7 @@ async def test_openai_compatible_custom_api_base(provider): print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize( @@ -5401,7 +5408,7 @@ async def test_openai_compatible_custom_api_video(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -7264,7 +7271,7 @@ def _openai_hallucinated_tool_call_mock_response( *args: object, **kwargs: object, ) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -7322,8 +7329,8 @@ def test_openai_hallucinated_tool_call(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=_openai_hallucinated_tool_call_mock_response, ) as mock_response: response = litellm.completion( @@ -7445,7 +7452,7 @@ def test_completion_novita_ai(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create" + openai_client, "post", return_value=MagicMock(spec=LegacyAPIResponse) ) as mock_call: mock_call.return_value.headers = {} mock_call.return_value.parse.return_value = litellm.ModelResponse( @@ -7464,11 +7471,11 @@ def test_completion_novita_ai(): # Verify model is passed correctly assert ( - mock_call.call_args.kwargs["model"] + mock_call.call_args.kwargs["body"]["model"] == "meta-llama/llama-3.3-70b-instruct" ) # Verify messages are passed correctly - assert mock_call.call_args.kwargs["messages"] == messages + assert mock_call.call_args.kwargs["body"]["messages"] == messages except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -7491,8 +7498,8 @@ def test_completion_novita_ai_dynamic_params(api_key): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=Exception("Invalid API key"), ) as mock_call: with pytest.raises(Exception, match="Invalid API key") as exc_info: @@ -7524,7 +7531,7 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): monkeypatch.setattr(litellm, "enable_preview_features", enable_preview_features) with patch.object( - client.chat.completions.with_raw_response, "create", return_value=MagicMock() + client, "post", return_value=MagicMock() ) as mock_completion: try: resp = litellm.completion( @@ -7538,11 +7545,11 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): mock_completion.assert_called_once() if enable_preview_features: - assert mock_completion.call_args.kwargs["metadata"] == { + assert mock_completion.call_args.kwargs["body"]["metadata"] == { "my-test-key": "my-test-value" } else: - assert "metadata" not in mock_completion.call_args.kwargs + assert "metadata" not in mock_completion.call_args.kwargs["body"] AZURE_TTS_BASE: Final = "https://tts.example.azure.com" 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_router/test_router_get_deployments.py b/tests/unit/test_router/test_router_get_deployments.py index 8b0c34f863c..07486762a64 100644 --- a/tests/unit/test_router/test_router_get_deployments.py +++ b/tests/unit/test_router/test_router_get_deployments.py @@ -5,6 +5,7 @@ import os import random import traceback from collections import defaultdict +from collections.abc import Callable from typing import Final, Literal import pytest @@ -110,9 +111,7 @@ def test_get_available_deployment_for_pass_through(): ) # Test that only pass-through deployment is returned - selected_model = router.get_available_deployment_for_pass_through( - "gpt-3.5-turbo" - ) + selected_model = router.get_available_deployment_for_pass_through("gpt-3.5-turbo") assert selected_model["litellm_params"]["model"] == "gpt-3.5-turbo" assert selected_model["litellm_params"]["use_in_pass_through"] is True @@ -259,9 +258,7 @@ def test_filter_pass_through_deployments(): healthy_deployments = router.get_model_list() # Filter pass-through deployments - pass_through_deployments = router._filter_pass_through_deployments( - healthy_deployments - ) + pass_through_deployments = router._filter_pass_through_deployments(healthy_deployments) # Should only have 2 deployments with use_in_pass_through=True assert len(pass_through_deployments) == 2 @@ -280,9 +277,7 @@ async def _weighted_async_deployment_ids(router: Router) -> tuple[str, ...]: random.seed(2025) deployments: Final = tuple( [ - await router.async_get_available_deployment( - model="shared", messages=None, request_kwargs={} - ) + await router.async_get_available_deployment(model="shared", messages=None, request_kwargs={}) for _ in range(1000) ] ) @@ -326,10 +321,7 @@ def test_weighted_selection_router( selected_ids: Final = ( asyncio.run(_weighted_async_deployment_ids(router)) if async_selection - else tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(1000) - ) + else tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(1000)) ) assert selected_ids.count("high") / len(selected_ids) > 0.89 @@ -355,10 +347,7 @@ def test_weighted_selection_router_no_rpm_set() -> None: }, ] ) - selected_ids: Final = tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(100) - ) + selected_ids: Final = tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(100)) assert set(selected_ids) == {"low", "high"} @@ -378,3 +367,61 @@ def test_model_group_aliases_select_the_target_group() -> None: assert deployment["model_name"] == "target" assert deployment["model_info"]["id"] == "target-deployment" + + +_ROUTABLE_MODEL_LIST: Final = [ + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-low", "api_key": "test-key"}, + "model_info": {"id": "low"}, + }, + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-high", "api_key": "test-key"}, + "model_info": {"id": "high"}, + }, + { + "model_name": "unrelated", + "litellm_params": {"model": "openai/unrelated", "api_key": "test-key"}, + "model_info": {"id": "unrelated"}, + }, +] + + +@pytest.mark.parametrize( + ("requested_model", "make_router", "expected_ids"), + ( + ("shared", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"low", "high"}), + ("high", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"high"}), + ("openai/unrelated", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"unrelated"}), + ( + "unknown", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, default_fallbacks=["shared"]), + {"low", "high"}, + ), + ( + "alias", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, model_group_alias={"alias": "openai/unrelated"}), + {"unrelated"}, + ), + ("unknown", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), set()), + ), + ids=( + "model-group", + "deployment-id", + "litellm-model-name", + "default-fallback", + "alias-to-a-litellm-model-name", + "unroutable", + ), +) +def test_get_routable_deployments_returns_every_deployment_routing_can_pick( + requested_model: str, make_router: Callable[[], Router], expected_ids: set[str] +) -> None: + router: Final = make_router() + + routable_ids: Final = { + deployment["model_info"]["id"] for deployment in router.get_routable_deployments(requested_model) + } + + assert routable_ids == expected_ids 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..423287e7c0f 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -231,9 +231,6 @@ "max-lines": { "count": 1 }, - "react-hooks/set-state-in-effect": { - "count": 1 - }, "react/no-unescaped-entities": { "count": 2 } @@ -1127,9 +1124,6 @@ "src/app/(dashboard)/usage/_components/components/UsagePageView.tsx": { "react-hooks/purity": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/app/(dashboard)/users/_components/BulkEditUsers.tsx": { @@ -1775,11 +1769,6 @@ "count": 1 } }, - "src/components/molecules/cost_optimization_feedback_banner.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/components/navbar.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2090,11 +2079,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 && (